diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 46e7db6..1be40c6 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -195,7 +195,7 @@ class FancyPtycho(CDIModel):
     @classmethod
     def from_dataset(cls,
                      dataset,
-                     probe_size=None,
+                     probe_shape=None,
                      randomize_ang=0,
                      n_modes=1,
                      n_obj_modes=1,
@@ -277,7 +277,7 @@ class FancyPtycho(CDIModel):
         )
 
         # Finally, initialize the probe and  object using this information
-        if probe_size is None:
+        if probe_shape is None:
             probe = tools.initializers.SHARP_style_probe(
                 dataset,
                 propagation_distance=propagation_distance,
@@ -288,7 +288,6 @@ class FancyPtycho(CDIModel):
                 dataset,
                 obj_basis,
                 probe_shape,
-                probe_size,
                 propagation_distance=propagation_distance,
             )
 
diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py
index 23df19a..afcd83e 100644
--- a/src/cdtools/models/multislice_ptycho.py
+++ b/src/cdtools/models/multislice_ptycho.py
@@ -181,7 +181,7 @@ class MultislicePtycho(CDIModel):
                      dataset,
                      dz,
                      nz,
-                     probe_size=None,
+                     probe_shape=None,
                      randomize_ang=0,
                      n_modes=1,
                      n_obj_modes=1,
@@ -262,7 +262,7 @@ class MultislicePtycho(CDIModel):
         )
 
         # Finally, initialize the probe and  object using this information
-        if probe_size is None:
+        if probe_shape is None:
             probe = tools.initializers.SHARP_style_probe(
                 dataset,
                 propagation_distance=propagation_distance,
@@ -273,7 +273,6 @@ class MultislicePtycho(CDIModel):
                 dataset,
                 obj_basis,
                 probe_shape,
-                probe_size,
                 propagation_distance=propagation_distance,
             )
 
diff --git a/src/cdtools/tools/initializers/initializers.py b/src/cdtools/tools/initializers/initializers.py
index 7e727af..bdb91b9 100644
--- a/src/cdtools/tools/initializers/initializers.py
+++ b/src/cdtools/tools/initializers/initializers.py
@@ -234,7 +234,7 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0, polariz
         polarizer = dataset.polarizer.tolist()
         analyzer = dataset.analyzer.tolist()
         factors = [(math.cos(math.radians(polarizer[idx] - analyzer[idx])))**2 for idx in range(len(dataset)) if (abs(polarizer[idx] - analyzer[idx]) > 5)]
-        avg_intensities = [t.sum(dataset[idx][1]) / factor[idx] for idx in range(len(dataset))]
+        avg_intensities = [t.sum(dataset[idx][1]) / factors[idx] for idx in range(len(dataset))]
 
     avg_intensity = t.mean(t.tensor(avg_intensities))
     probe_intensity = t.sum(t.abs(probe)**2)