Update train.py

train.py

@@ -27,14 +27,14 @@ def parse_args():
 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)
+def get_dataloaders(data_path, transforms, batch_size, fine_tuning):
+    train_dataset = TrainDataset(data_path, transforms)
     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
 
 def get_models(device):