Update is_floating_point to handle bfloat16 (apache#7133)

* Add div_ and is_floating_point operators * Add handling of exprs to op, update tests * Properly handle bfloat16 in is_floating_point * Revert test changes * revert whitespace changes
electriclilies · Feb 18, 2021 · a915b63 · a915b63
1 parent 4dc459a
commit a915b63
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -2067,7 +2067,7 @@ def is_floating_point(self, inputs, input_types):
         else:
             input_type = input_types[0]
 
-        is_float = input_type in ["float32", "float64", "float16"]
+        is_float = input_type in ["float32", "float64", "float16", "bfloat16"]
         return _expr.const(is_float)
 
     # Operator mappings