device

Lightning-AI · May 20, 2020 · 405fcf1 · 405fcf1
1 parent b022f0d
commit 405fcf1
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tests/core/test_memory.py b/tests/core/test_memory.py
@@ -61,8 +61,13 @@ def forward(self, x, y):
     assert model.training
 
 
-def test_rnn_summary_shapes():
-    model = ParityRNN()
+@pytest.mark.parametrize('device', [
+    torch.device('cpu'),
+    torch.device('cuda', 0)
+])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU.")
+def test_rnn_summary_shapes(device):
+    model = ParityRNN().to(device)
 
     b = 3
     t = 5