diff --git a/src/neuron_proofreader/merge_proofreading/search_datasets.py b/src/neuron_proofreader/merge_proofreading/search_datasets.py
index d2666da..9cf865d 100644
--- a/src/neuron_proofreader/merge_proofreading/search_datasets.py
+++ b/src/neuron_proofreader/merge_proofreading/search_datasets.py
@@ -283,7 +283,6 @@ def __init__(
         is_multimodal=False,
         min_search_size=0,
         prefetch=64,
-        step_size=10,
         subgraph_radius=100,
     ):
         # Call parent class
diff --git a/src/neuron_proofreader/models/new_vision_models.py b/src/neuron_proofreader/models/new_vision_models.py
index 78327e5..5fe88f0 100644
--- a/src/neuron_proofreader/models/new_vision_models.py
+++ b/src/neuron_proofreader/models/new_vision_models.py
@@ -38,6 +38,12 @@ def __init__(
 
         Parameters
         ----------
+        input_shape : Tuple[int]
+            Input shape.
+        base_channels : int, optional
+            Number of channels in first convolutional layer. Default is 16.
+        channel_multiplier : float, optional
+            ...
         output_dim : int, optional
             Dimension of output. Default is 1.
         dropout : float, optional