Update train.py

train.py

@@ -29,15 +29,18 @@ def get_transforms():
 
 def get_dataloaders(data_path, transforms, batch_size, fine_tuning, mult_number):
     train_dataset = TrainDataset(data_path, transforms, mults_amount=mult_number)
+    print("Train Dataset Length:", len(train_dataset))  # Debug print
     train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
     
     if fine_tuning:
         finetuning_dataset = FineTuningDataset(data_path, transforms, mult_amount=mult_number)
+        print("FineTuning Dataset Length:", len(finetuning_dataset))  # Debug print
         finetuning_dataloader = torch.utils.data.DataLoader(finetuning_dataset, batch_size=batch_size, shuffle=True)
         return train_dataloader, finetuning_dataloader
 
     return train_dataloader, None
 
+
 def get_models(device):
     generator = Generator()
     extractor = get_seresnext_extractor()