Merge remote-tracking branch 'origin/feature/background_estimator' in…

…to feature/background_estimator # Conflicts: # lensless/utils/dataset.py
LCAV · Dec 17, 2024 · f5a6bf6 · f5a6bf6
2 parents 2f40b1b + 18e8e66
commit f5a6bf6
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/scripts/recon/background_estimator.py b/scripts/recon/background_estimator.py
@@ -413,9 +413,9 @@ def log_predictions(self, inputs, outputs, background, epoch):
         background_np = background.cpu().numpy()
 
         # Remove batch dimension
-        inputs_np = inputs_np.squeeze(0)  # Shape: [C, H, W]
-        outputs_np = outputs_np.squeeze(0)  # Shape: [C, H, W]
-        background_np = background_np.squeeze(0)  # Shape: [C, H, W]
+        inputs_np = inputs_np[0] # Shape: [C, H, W]
+        outputs_np = outputs_np[0] # Shape: [C, H, W]
+        background_np = background_np[0] # Shape: [C, H, W]
 
         # Remove the extra channel from inputs
         if inputs_np.shape[0] > 3: