Update simulation.py for autograd (#11)

Inplace operation /= and *= break autograd
ebezzam · Aug 17, 2023 · 82dfb08 · 82dfb08
1 parent c07863a
commit 82dfb08
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/waveprop/simulation.py b/waveprop/simulation.py
@@ -163,8 +163,8 @@ def propagate(self, obj, return_object_plane=False):
                 image_plane = add_shot_noise(image_plane, snr_db=self.snr_db)
 
             # 5) Quantize as on sensor
-            image_plane /= image_plane.max()
-            image_plane *= self.max_val
+            image_plane = image_plane / image_plane.max()
+            image_plane = image_plane * self.max_val
             if torch.is_tensor(image_plane):
                 image_plane = image_plane.to(self.output_dtype)
             else: