pytorch
1fbc08c7 - Add Autocast support for Einsum (#71916)

Commit
2 years ago
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)
Author
Thiago Crepaldi
Committer
Parents
Loading