Clarify doc string

skorch-dev · Jan 8, 2025 · e17b718 · e17b718
1 parent d740ff9
commit e17b718
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/skorch/callbacks/lr_scheduler.py b/skorch/callbacks/lr_scheduler.py
@@ -87,10 +87,12 @@ def simulate(self, steps, initial_lr, step_args=None):
         initial_lr: float
           Initial learning rate
 
-        step_args: Any
-          Argument to the `.step()` function of the policy. If it is an
+        step_args: None or float or List[float] (default=None)
+          Argument to the ``.step()`` function of the policy. If it is an
           indexable object the simulation will try to associate every step of
-          the simulation with an entry in ``step_args``.
+          the simulation with an entry in ``step_args``. Scalar values are
+          passed at every step, unchanged. In the default setting (``None``)
+          no additional arguments are passed to ``.step()``.
 
         Returns
         -------