amend

pytorch · Dec 6, 2023 · 91fbde4 · 91fbde4
1 parent e1c38f8
commit 91fbde4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py
@@ -273,7 +273,8 @@ def _lstm(self, x, hx):
                     x_t = F.dropout(_h_t, p=self.dropout, training=self.training)
                 else:  # No dropout after the last layer
                     x_t = _h_t
-
+            h_t = h_t_out
+            c_t = c_t_out
             outputs.append(x_t)
 
         outputs = torch.stack(outputs, dim=1)