diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 7825c43f..88734e34 100644
--- 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
index 0e8c9456..11ad4c60 100644
--- 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,
             )