Fix a bug in set float32 precision (#2271)

Author: Nicorgi

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
         )

torchtune/training/precision.py

@@ -34,11 +34,15 @@ 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
+<<<<<<< HEAD
 <<<<<<< HEAD
     if not (torch.cuda.is_available() or is_npu_available):
 =======
     if not torch.cuda.is_available() or not is_npu_available:
 >>>>>>> c5f20b96 (Add Ascend NPU as a backend for single device recipes (#2234))
+=======
+    if not (torch.cuda.is_available() or is_npu_available):
+>>>>>>> 890deab3 (Fix a bug in set float32 precision (#2271))
         return
     # set precision for matrix multiplications
     torch.set_float32_matmul_precision(precision)