Update model.py

src/models/model.py

@@ -390,6 +390,7 @@ class Summarization:
 
         mlflow.pytorch.autolog(log_models=False)
 
+        with mlflow.start_run() as run:  # noqa: F841
             trainer.fit(self.T5Model, self.data_module)
 
     def load_model(