Support BF16 in TorchTNT logger

Reviewed By: galrotem

Differential Revision: D70537893

Author: JKSenthil

tests/utils/loggers/test_utils.py

@@ -29,3 +29,8 @@ def test_scalar_to_float(self) -> None:
 
         valid_ndarray = np.array([[[float_x]]])
         self.assertAlmostEqual(scalar_to_float(valid_ndarray), float_x)
+
+    def test_scalar_to_float_bf16(self) -> None:
+        float_x = 3.45
+        valid_tensor = torch.Tensor([float_x]).to(torch.bfloat16)
+        self.assertAlmostEqual(scalar_to_float(valid_tensor), float_x, delta=0.01)

torchtnt/utils/loggers/utils.py

@@ -20,7 +20,7 @@ def scalar_to_float(scalar: Scalar) -> float:
                 f"Scalar tensor must contain a single item, {numel} given."
             )
 
-        return float(scalar.cpu().detach().numpy().item())
+        return float(scalar.cpu().detach().float().numpy().item())
     elif isinstance(scalar, ndarray):
         numel = scalar.size
         if numel != 1: