fix error

jiexiaou · Sep 21, 2022 · 63a8c63 · 63a8c63
1 parent a7625ac
commit 63a8c63
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/dataset/imgdata.py b/dataset/imgdata.py
@@ -165,6 +165,9 @@ def __getitem__(self, index):
         tar_img = tar_img[:, rr:rr + self.crop_size, cc:cc + self.crop_size]
         # Data Augmentations
         aug = random.randint(0, 3)
+        # Crop patch
+        inp_img = inp_img[:, rr:rr + self.crop_size, cc:cc + self.crop_size]
+        tar_img = tar_img[:, rr:rr + self.crop_size, cc:cc + self.crop_size]
         if aug == 1:
             inp_img = inp_img.flip(1)
             tar_img = tar_img.flip(1)