From 04fa69c31a73b5fbb96251ccf19c4b47cd188c70 Mon Sep 17 00:00:00 2001
From: Ilyes Batatia <48651863+ilyes319@users.noreply.github.com>
Date: Thu, 21 Nov 2024 14:22:10 +0000
Subject: [PATCH] add back pytest device

---
 tests/test_cueq.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_cueq.py b/tests/test_cueq.py
index 52ed1e3c..21ec0e2e 100644
--- 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",
         [