amend

facebookresearch · Aug 1, 2024 · a1191dc · a1191dc
1 parent 60aeb3e
commit a1191dc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/benchmarl/models/gru.py b/benchmarl/models/gru.py
@@ -250,7 +250,7 @@ def run_net(self, input, is_init, h_0):
                 output, h_n = self.vmap_func_module(
                     self._empty_gru,
                     (0, -2, -2, -3),
-                    (-2, -2),
+                    (-2, -3),
                 )(self.params, input, is_init, h_0)
         else:
             with self.params.to_module(self._empty_gru):