From a1191dcc0ac5ec4ed458d19a520d9ccfb01d403b Mon Sep 17 00:00:00 2001
From: Matteo Bettini <mb2389@cl.cam.ac.uk>
Date: Thu, 1 Aug 2024 08:24:34 +0200
Subject: [PATCH] amend

---
 benchmarl/models/gru.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmarl/models/gru.py b/benchmarl/models/gru.py
index 82510536..d4ebe062 100644
--- 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):