Fix psf shape for torch array.

LCAV · May 14, 2024 · aca35f1 · aca35f1
1 parent 754502e
commit aca35f1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -23,7 +23,7 @@ Changed
 Bugfix
 ~~~~~~
 
-- Nothing
+- Wavelength input to ``compute_psf`` for Mask class.
 
 
 1.0.7 - (2024-05-14)

diff --git a/lensless/hardware/mask.py b/lensless/hardware/mask.py
@@ -219,7 +219,7 @@ def compute_psf(self, distance_sensor=None, wavelength=None, intensity=True):
 
         if self.use_torch:
             psf = torch.zeros(
-                tuple(self.resolution) + (len(self.psf_wavelength),),
+                tuple(self.resolution) + (len(wavelength),),
                 dtype=torch.complex64,
                 device=self.torch_device,
             )