diff --git a/datasets/fake/generate_dataset.py b/datasets/fake/generate_dataset.py
index 30c7aa1..929bf03 100644
--- a/datasets/fake/generate_dataset.py
+++ b/datasets/fake/generate_dataset.py
@@ -18,10 +18,12 @@ NUMBER_OF_IMAGES_PER_PHONE = 10_000
 # Compared to images being 1.
 PRNU_FACTOR = 0.1
 
+IMAGE_SIZE_SHAPE = (IMAGE_SIZE, IMAGE_SIZE)
+
 # Generate PRNUs and images of phones.
 # Is such `np.maximum` probabilistically correct with our theoretical method? See #19.
 def randomImage(scale):
-    return np.random.normal(loc = 0, scale = scale, size = (IMAGE_SIZE, IMAGE_SIZE))
+    return np.random.normal(loc = 0, scale = scale, size = IMAGE_SIZE_SHAPE)
 
 imagesWithoutPrnu = [[randomImage(scale = 1) for _ in range(NUMBER_OF_IMAGES_PER_PHONE)] for phoneIndex in range(NUMBER_OF_PHONES)]
 
@@ -47,11 +49,11 @@ plt.title('RMS between actual PRNU and the mean of the first $N$ images with PRN
 plt.xlabel('$N$ first images with PRNU')
 plt.ylabel('RMS')
 rmss = []
-mean = 0
+mean = np.zeros(IMAGE_SIZE_SHAPE)
 for imageIndex in range(NUMBER_OF_IMAGES_PER_PHONE):
-    rms = rmsDiffNumpy(imagesWithPrnu[0][imageIndex], prnus[0])
-    mean = ((mean * imageIndex) + rms) / (imageIndex + 1)
-    rmss += [mean]
+    mean = (mean * imageIndex + imagesWithPrnu[0][imageIndex]) / (imageIndex + 1)
+    rms = rmsDiffNumpy(mean, prnus[0])
+    rmss += [rms]
 plt.plot(rmss)
 plt.show()