Set device as CPU in failing tests

Rose-STL-Lab · Apr 29, 2024 · f254258 · f254258
1 parent f6d9c97
commit f254258
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -37,8 +37,6 @@ jobs:
 
       - name: Run tests
         run: poetry run pytest tests/ --cov=torchts
-        env:
-          PYTORCH_MPS_HIGH_WATERMARK_RATIO: 0.0
 
       - name: Generate coverage report
         run: poetry run coverage xml

diff --git a/tests/nn/test_ode.py b/tests/nn/test_ode.py
@@ -99,6 +99,7 @@ def test_step_backward(euler_model):
 def test_fit(euler_model):
     """Test the step and backward function"""
     torch.manual_seed(0)
+    torch.set_default_device("cpu")
     model, _ = euler_model
     model.fit(torch.Tensor([[1.0]]), torch.Tensor([[1.1]]), max_epochs=1, batch_size=1)
     coeffs = model.get_coeffs()

diff --git a/tests/test_model.py b/tests/test_model.py
@@ -59,6 +59,7 @@ def test_step(mocker):
 
 def test_train():
     torch.manual_seed(0)
+    torch.set_default_device("cpu")
 
     slope_init = 2
     intercept_init = -1