fix for numpy>=2.0

ahalev · Aug 8, 2024 · e22b7ee · e22b7ee
1 parent fc2c64f
commit e22b7ee
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/pymgrid/modules/base/base_module.py b/src/pymgrid/modules/base/base_module.py
@@ -149,9 +149,10 @@ def step(self, action, normalized=True):
         except (IndexError, TypeError):
             if not isinstance(denormalized_action, (float, int)):
                 try:
-                    flat_dim = np.product(denormalized_action.shape)
-                    assert flat_dim == 0
-                except (AttributeError, AssertionError):
+                    flat_dim = np.prod(denormalized_action.shape)
+                    if flat_dim != 0:
+                        raise ValueError(f'Bad action {denormalized_action}')
+                except AttributeError:
                     raise ValueError(f'Bad action {denormalized_action}')
                 else:
                     denormalized_action = 0.0