update torch compile

michaelfeil · Jan 3, 2025 · 6ddba7b · 6ddba7b
1 parent 9544138
commit 6ddba7b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/libs/infinity_emb/tests/unit_test/transformer/crossencoder/test_torch_crossencoder.py b/libs/infinity_emb/tests/unit_test/transformer/crossencoder/test_torch_crossencoder.py
@@ -11,7 +11,7 @@
 
 device = Device.cpu if torch.backends.mps.is_available() else Device.auto
 
-SHOULD_TORCH_COMPILE = sys.platform == "linux" and sys.version_info < (3, 12)
+SHOULD_TORCH_COMPILE = sys.platform == "linux" and sys.version_info < (3, 12) and torch.cuda.is_available()
 
 
 def test_crossencoder():
@@ -43,7 +43,7 @@ def test_crossencoder():
 def test_patched_crossencoder_vs_sentence_transformers():
     model = CrossEncoderPatched(
         engine_args=EngineArgs(
-            model_name_or_path="mixedbread-ai/mxbai-rerank-xsmall-v1", compile=True, device=device
+            model_name_or_path="mixedbread-ai/mxbai-rerank-xsmall-v1", compile=SHOULD_TORCH_COMPILE, device=device
         )
     )
     model_unpatched = CrossEncoder("mixedbread-ai/mxbai-rerank-xsmall-v1")