Update train.py

train.py

@@ -28,14 +28,15 @@ def get_transforms():
     return albu.Compose([albu.RandomCrop(512, 512, always_apply = True), albu.HorizontalFlip(p = 0.5)], p = 1.)
 
 def get_dataloaders(data_path, transforms, batch_size, fine_tuning, mult_number):
-    train_dataset = TrainDataset(data_path, transforms, mult_number)
-    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size = batch_size, shuffle = True)
+    train_dataset = TrainDataset(data_path, transforms, mults_amount=mult_number)
+    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
     
     if fine_tuning:
-        finetuning_dataset = FineTuningDataset(data_path, transforms)
-        finetuning_dataloader = torch.utils.data.DataLoader(finetuning_dataset, batch_size = batch_size, shuffle = True)
-    
-    return train_dataloader, finetuning_dataloader
+        finetuning_dataset = FineTuningDataset(data_path, transforms, mult_amount=mult_number)
+        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()