fixes

src/models/model.py

@@ -7,7 +7,7 @@ from transformers import (
 )
 from torch.utils.data import Dataset, DataLoader
 import pytorch_lightning as pl
-from pytorch_lightning.loggers import MLFlowLogger
+from pytorch_lightning.loggers import MLFlowLogger, WandbLogger
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
 from pytorch_lightning import LightningDataModule
@@ -327,11 +327,10 @@ class Summarization:
             learning_rate=learning_rate, adam_epsilon=adam_epsilon, weight_decay=weight_decay
         )
 
-        MLlogger = MLFlowLogger(experiment_name="Summarization",
-                                tracking_uri="https://dagshub.com/gagan3012/summarization.mlflow")
-                                #save_dir="reports/training_metrics.txt"
+        # MLlogger = MLFlowLogger(experiment_name="Summarization",
+         #                       tracking_uri="https://dagshub.com/gagan3012/summarization.mlflow")
 
-        # WandLogger = WandbLogger(project="summarization-dagshub")
+        WandLogger = WandbLogger(project="summarization-dagshub")
 
         # logger = DAGsHubLogger(metrics_path='reports/training_metrics.txt')
 
@@ -352,7 +351,7 @@ class Summarization:
         gpus = -1 if use_gpu and torch.cuda.is_available() else 0
 
         trainer = Trainer(
-            logger=MLlogger,
+            logger=WandLogger,
             callbacks=early_stop_callback,
             max_epochs=max_epochs,
             gpus=gpus,