fix bug in set_float32_precision

fix bug in UT
pytorch · Jan 16, 2025 · 969c351 · 969c351
1 parent 28c9ec2
commit 969c351
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/tests/torchtune/training/test_precision.py b/tests/torchtune/training/test_precision.py
@@ -57,8 +57,7 @@ def test_error_bf16_unsupported(self, mock_verify):
             get_dtype(torch.bfloat16)
 
     @pytest.mark.skipif(not cuda_available, reason="The test requires GPUs to run.")
-    @mock.patch("torchtune.training.precision.is_npu_available", return_value=True)
-    def test_set_float32_precision(self, mock_npu_available) -> None:
+    def test_set_float32_precision(self) -> None:
         setattr(  # noqa: B010
             torch.backends, "__allow_nonbracketed_mutation_flag", True
         )

diff --git a/torchtune/training/precision.py b/torchtune/training/precision.py
@@ -34,7 +34,7 @@ def _set_float32_precision(precision: str = "high") -> None:
         precision (str): The setting to determine which datatypes to use for matrix multiplication and convolution operations.
     """
     # Not relevant for non-CUDA or non-NPU devices
-    if not torch.cuda.is_available() or not is_npu_available:
+    if not (torch.cuda.is_available() or is_npu_available):
         return
     # set precision for matrix multiplications
     torch.set_float32_matmul_precision(precision)