Merge pull request #81 from tsugumi-sys/fix/failed-tests-device-error

fix failed tests on gpu device
tsugumi-sys · Jan 15, 2024 · c2d533a · c2d533a
2 parents a6a8f28 + 78ac713
commit c2d533a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tests/self_attention_memory_convlstm/test_model.py b/tests/self_attention_memory_convlstm/test_model.py
@@ -18,6 +18,6 @@ def test_SAMConvLSTM():
         },
     }
 
-    model = SAMConvLSTM(**model_params)
+    model = SAMConvLSTM(**model_params).to(DEVICE)
     output = model(torch.rand((2, 1, 3, 8, 8), dtype=torch.float, device=DEVICE))
     assert output.size() == (2, 1, 3, 8, 8)
diff --git a/tests/test_model/test_model.py b/tests/test_model/test_model.py
@@ -3,6 +3,7 @@
 import pytest
 import torch
 
+from core.constants import DEVICE
 from tests.test_model.model import TestModel
 
 
@@ -14,6 +15,6 @@ def test_TestModel_successfully_build(
     return_sequences: bool, expected_output_size: Tuple
 ):
     # return_sequence
-    model = TestModel(return_sequences)
-    output = model(torch.rand((5, 6, 3, 16, 16), dtype=torch.float))
+    model = TestModel(return_sequences).to(DEVICE)
+    output = model(torch.rand((5, 6, 3, 16, 16), dtype=torch.float, device=DEVICE))
     assert output.size() == expected_output_size