Remove changes in training.

LCAV · Aug 22, 2024 · 38a6f0c · 38a6f0c
1 parent f34bde2
commit 38a6f0c
Showing 1 changed file with 6 additions and 19 deletions.
diff --git a/lensless/recon/utils.py b/lensless/recon/utils.py
@@ -873,28 +873,15 @@ def train_epoch(self, data_loader):
         for batch in pbar:
 
             # get batch
-            flip_lr = None
-            flip_ud = None
-            X = batch[0].to(self.device)
-            y = batch[1].to(self.device)
-            if self.train_multimask or self.train_random_flip:
-                psfs = batch[2]
+            if self.train_random_flip:
+                X, y, psfs, flip_lr, flip_ud = batch
+                psfs = psfs.to(self.device)
+            elif self.train_multimask:
+                X, y, psfs = batch
                 psfs = psfs.to(self.device)
             else:
+                X, y = batch
                 psfs = None
-            if self.train_random_flip:
-                flip_lr = batch[3]
-                flip_ud = batch[4]
-
-            # if self.train_random_flip:
-            #     X, y, psfs, flip_lr, flip_ud = batch
-            #     psfs = psfs.to(self.device)
-            # elif self.train_multimask:
-            #     X, y, psfs = batch
-            #     psfs = psfs.to(self.device)
-            # else:
-            #     X, y = batch
-            #     psfs = None
 
             random_rotate = False
             if self.random_rotate: