add back pytest device

ACEsuit · Nov 21, 2024 · 04fa69c · 04fa69c
1 parent 280dfe0
commit 04fa69c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/test_cueq.py b/tests/test_cueq.py
@@ -74,6 +74,7 @@ def batch(self, device: str):
         batch = next(iter(data_loader))
         return batch.to(device).to_dict()
 
+    @pytest.mark.parametrize("device", ["cpu"])
     @pytest.mark.parametrize(
         "interaction_cls_first",
         [