Update dataset/datasets.py

dataset/datasets.py

@@ -43,10 +43,12 @@ class TrainDataset(torch.utils.data.Dataset):
 
         return bw_img, color_img, hint
 
-class FineTuningDataset(Dataset):
-    def __init__(self, data_path, transform=None):
+class FineTuningDataset(torch.utils.data.Dataset):
+    def __init__(self, data_path, transform=None, mult_amount=1):
+        self.data = [x for x in os.listdir(os.path.join(data_path, 'real_manga')) if x.endswith('_dfm.png')]
         self.data_path = data_path
         self.transform = transform
+        self.mult_amount = mult_amount
         self.ToTensor = transforms.ToTensor()
 
         self.file_list = [x for x in os.listdir(os.path.join(data_path, 'real_manga')) if x.endswith('_dfm.png')]