amend

pytorch · Dec 6, 2023 · b2d37e6 · b2d37e6
1 parent ac6618d
commit b2d37e6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py
@@ -220,7 +220,7 @@ def _lstm_cell(x, hx, cx, weight_ih, bias_ih, weight_hh, bias_hh):
 
         # gates = F.linear(x, weight_ih, bias_ih) + F.linear(hx, weight_hh, bias_hh)
         if bias_ih is not None:
-            gates = x @ weight_ih.T + bias_ih + hx @ weight_ih.T + bias_hh
+            gates = x @ weight_ih.T + bias_ih + hx @ weight_hh.T + bias_hh
         else:
             gates = x @ weight_ih.T + hx @ weight_hh.T