Update transformer_engine/pytorch/tensor/quantized_tensor.py

Co-authored-by: Tim Moon <[email protected]>
Signed-off-by: ZhiyiDanielSu <[email protected]>

Author: zobeideThePlayer

transformer_engine/pytorch/tensor/quantized_tensor.py

@@ -427,7 +427,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
                 dst.quantize_(src)
             else:
                 if isinstance(src, QuantizedTensor):
-                    src = src.dequantize(dtype=dst.dtype)
+                    dtype = dst.dtype
+                    if dtype not in (torch.float32, torch.float16, torch.bfloat16):
+                        dtype = torch.float32
+                    src = src.dequantize(dtype=dtype)
                 dst.copy_(src)
             return None