diff --git a/src/pymgrid/microgrid/microgrid.py b/src/pymgrid/microgrid/microgrid.py
index 69899db3..12533173 100644
--- a/src/pymgrid/microgrid/microgrid.py
+++ b/src/pymgrid/microgrid/microgrid.py
@@ -237,7 +237,7 @@ def step(self, control, normalized=True):
 
         Parameters
         ----------
-        control : dict[str, list[float]]
+        control : dict[str, np.typing.arrayLike]
             Actions to pass to each fixed module.
         normalized : bool, default True
             Whether ``control`` is a normalized value or not. If not, each module de-normalizes its respective action.
diff --git a/src/pymgrid/modules/grid_module.py b/src/pymgrid/modules/grid_module.py
index 4bc0ccaa..ec224fc9 100644
--- a/src/pymgrid/modules/grid_module.py
+++ b/src/pymgrid/modules/grid_module.py
@@ -26,7 +26,7 @@ class GridModule(BaseTimeSeriesMicrogridModule):
         If n_features=4, time series of ``(import_price, export_price, co2_per_kwH, grid_status)``
         in each column, respectively. ``time_series[:, -1]`` -- the grid status -- must be binary.
 
-    forecaster : callable, float, "oracle", or None, default None.
+    forecaster : callable, float, "oracle", None, default None.
         Function that gives a forecast n-steps ahead.
 
         * If ``callable``, must take as arguments ``(val_c: float, val_{c+n}: float, n: int)``, where