Don't convert bfloat16 to float32 (#41)

tenstorrent · Oct 30, 2024 · c058541 · c058541
1 parent 69089f2
commit c058541
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -5339,8 +5339,10 @@ def _convert_data_type(input_type, default_dtype=None):
     input_type = input_type.lower()
     if input_type in ["double", "float64", "torch.float64"]:
         return "float64"
-    elif input_type in ["float", "float32", "torch.float32", "torch.bfloat16", "bfloat16"]:
+    elif input_type in ["float", "float32", "torch.float32"]:
         return "float32"
+    elif input_type in ["torch.bfloat16", "bfloat16"]:
+        return "bfloat16"
     elif input_type in ["half", "float16", "torch.float16"]:
         return "float16"
     elif input_type in ["long", "int64", "torch.int64"]: