Add Autocast support for Einsum (#71916)
Summary:
ONNX spec for Einsum requires all inputs to be the same dtype.
PyTorch runtime does not allow executing aten::einsum with
mismatching types by default, so the export would never succeed,
However, when the model is wrapped by `torch.autocast()`,
the run succeeds and the ONNX converter will create an Einsum ONNX node
with mismatch types as input, which is not allowed by the aforementioned schema.
This PR adds onnx::Einsum to the Autocast enabled list, which outputs lower resolution tensors
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71916
Reviewed By: ngimel
Differential Revision: D34629666
Pulled By: malfet
fbshipit-source-id: ec757bb87190a5b7512969e10a32450e9e1f87a1
(cherry picked from commit 7f2b5a6408ae34a6b9f858c3e9f5970b64ca1b4b)