Update dataset/definition_dataset.py

dataset/definition_dataset.py

@@ -29,7 +29,7 @@ class SemanticSegmentationDatasetFusion(Dataset):
     - transforms applied to images (ToTensor + Normalize) and mask -> LongTensor
     - optional JOINT augmentations applied consistently to HR/LR + masks
     """
-    def __init__(self, image_dir: Path, mask_dir: Path, transform: Optional[transforms.Compose] = None):
+    def __init__(self, image_dir: Path, mask_dir: Path, transform: Optional[transforms.Compose] = None, augment = None):
         self.image_dir = Path(image_dir)
         self.mask_dir = Path(mask_dir)
         self.image_filenames = sorted(os.listdir(self.image_dir))