minor changes

img_dataset.py

@@ -11,8 +11,8 @@ property_label_to_idx = {'淡红': 0, '淡白': 1, '红': 2, '淡紫': 3, '紫
 
 class Image_DataSet(torch.utils.data.Dataset):
 
-    def __init__(self, transform=None):
-        self.processed_dir = './processed-data'
+    def __init__(self, processed_dir = './processed-data', transform=None):
+        self.processed_dir = processed_dir
         self.transform = transform
 
     def __getitem__(self, index):

train_tongue_diagnosis.py

@@ -58,8 +58,8 @@ transform_train = transforms.Compose([
     transforms.ToTensor(),
     normalizer
 ])
-
-tongue_dataset_train = Image_DataSet(transform=transform_train)
+processed_dir = './processed-data'
+tongue_dataset_train = Image_DataSet(processed_dir=processed_dir, transform=transform_train)
 tongue_train_loader = torch.utils.data.DataLoader(tongue_dataset_train, batch_size=args.batch_size, shuffle=True,
                                                   **kwargs)