Add embedding cast in third_party/tvm/python/tvm/relay/frontend/pytor…

…ch.py to cast weight to bfloat16 if they are float32. (#55)

python/tvm/relay/frontend/pytorch.py

@@ -2700,7 +2700,10 @@ def embedding(self, inputs, input_types):
             #  exposes a few bugs in tt-mlir https://github.com/tenstorrent/tt-mlir/issues/1215
             logger.warning("Casting input indices of embedding op from {} to int32", indicies_dtype)
             indices = tvm.relay.cast(indices, "int32")
-        return _op.embedding(weight, indices, axis=0)
+        # cast the weight to bfloat16 if it is float32
+        if weight.type_annotation.dtype == "float32":
+            weight = tvm.relay.cast(weight, "bfloat16")
+        return tvm.relay.cast(_op.embedding(weight, indices, axis=0), "float32")
 
     def embedding_bag(self, inputs, input_types):