Fix bf16 check in skipped test

Signed-off-by: Alex-Brooks <[email protected]>
foundation-model-stack · Mar 21, 2024 · efa446f · efa446f
1 parent 05663e8
commit efa446f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -485,7 +485,8 @@ def test_no_packing_needs_reponse_template():
 
 ### Tests for model dtype edge cases
 @pytest.mark.skipif(
-    torch.cuda.is_bf16_supported(), reason="Only runs if bf16 is unsupported"
+    not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
+    reason="Only runs if bf16 is unsupported",
 )
 def test_bf16_still_tunes_if_unsupported():
     """Ensure that even if bf16 is not supported, tuning still works without problems."""