diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py
index 9d09b1ce..81b3f886 100644
--- a/autoarray/dataset/imaging/dataset.py
+++ b/autoarray/dataset/imaging/dataset.py
@@ -400,7 +400,7 @@ def apply_noise_scaling(
             data = np.where(np.invert(mask), 0.0, self.data.native)
         elif data_noise_sigma is not None:
             random_noise = np.random.normal(
-                loc=data_noise_mean, scale=data_noise_sigma, size=data.shape_native
+                loc=data_noise_mean, scale=data_noise_sigma, size=self.data.shape_native
             )
             data = np.where(mask_2d, random_noise, data.native)