Fix support for bf16 in to_numpy

BloodAxe · Oct 14, 2024 · fcae22a · fcae22a
1 parent 5b65fa1
commit fcae22a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pytorch_toolbelt/utils/torch_utils.py b/pytorch_toolbelt/utils/torch_utils.py
@@ -153,6 +153,8 @@ def to_numpy(x: Union[torch.Tensor, np.ndarray, Any, None]) -> Union[np.ndarray,
     """
     if x is None:
         return None
+    elif torch.is_tensor(x) and x.dtype == torch.bfloat16:
+        return x.data.float().cpu().numpy()
     elif torch.is_tensor(x):
         return x.data.cpu().numpy()
     elif isinstance(x, np.ndarray):