Update test/test_modules.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
pytorch · Dec 1, 2023 · 3dd0a52 · 3dd0a52
1 parent 5331647
commit 3dd0a52
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/test_modules.py b/test/test_modules.py
@@ -1297,8 +1297,8 @@ def test_python_lstm(device, bias, dropout, batch_first, num_layers):
     else:
         input = torch.randn(T, B, 10, device=device)
 
-    h0 = torch.randn(num_layers, 5, 20).to(device)
-    c0 = torch.randn(num_layers, 5, 20).to(device)
+    h0 = torch.randn(num_layers, 5, 20, device=device)
+    c0 = torch.randn(num_layers, 5, 20, device=device)
 
     # Test without hidden states
     with torch.no_grad():