Update test/test_modules.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
pytorch · Dec 1, 2023 · 6c0fff1 · 6c0fff1
1 parent e3012f2
commit 6c0fff1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/test/test_modules.py b/test/test_modules.py
@@ -1197,6 +1197,8 @@ def test_python_lstm_cell(device, bias):
     lstm_cell1 = LSTMCell(10, 20, device=device, bias=bias)
     lstm_cell2 = nn.LSTMCell(10, 20, device=device, bias=bias)
 
+    lstm_cell1.load_state_dict(lstm_cell2.state_dict())
+
     # Make sure parameters match
     for (k1, v1), (k2, v2) in zip(
         lstm_cell1.named_parameters(), lstm_cell2.named_parameters()