jnp.float16 -> torch.float16

pytorch · Apr 5, 2024 · b08a9ad · b08a9ad
1 parent 40ae313
commit b08a9ad
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/experimental/torch_xla2/torch_xla2/tensor.py b/experimental/torch_xla2/torch_xla2/tensor.py
@@ -115,6 +115,7 @@ def t2j_dtype(dtype):
 
 def j2t_dtype(dtype):
   return {
+      jnp.float16: torch.float16,
       jnp.bfloat16: torch.bfloat16,
       jnp.double: torch.double,
       jnp.float32: torch.float32,