[CI] revert device in test_export_static_cache (#39662)

* revert device * add todo
2025-07-25 16:36:12 +01:00
parent 850bdeaa95
commit 5d0ba3e479
10 changed files with 10 additions and 10 deletions
--- a/tests/models/gemma2/test_modeling_gemma2.py
+++ b/tests/models/gemma2/test_modeling_gemma2.py
@@ -335,7 +335,7 @@ class Gemma2IntegrationTest(unittest.TestCase):
        ].shape[-1]

        # Load model
-        device = torch_device
+        device = "cpu"  # TODO (joao / export experts): should be on `torch_device`, but causes GPU OOM
        dtype = torch.bfloat16
        cache_implementation = "static"
        attn_implementation = "sdpa"