diff --git a/src/scgenerator/physics/units.py b/src/scgenerator/physics/units.py
index 7a13bb8..c4424ac 100644
--- a/src/scgenerator/physics/units.py
+++ b/src/scgenerator/physics/units.py
@@ -346,7 +346,7 @@ def to_WL(spectrum: np.ndarray, lambda_: np.ndarray) -> np.ndarray:
 
     Parameters
     ----------
-    spectrum : np.ndarray, shape (n, )
+    spectrum : np.ndarray, shape (..., n)
         intensity spectrum
     lambda_ : np.ndarray, shape (n, )
         wavelength in m