[torch.onnx.export] Fix onnx export on big endian machines#167816
[torch.onnx.export] Fix onnx export on big endian machines#167816tungld wants to merge 2 commits intopytorch:mainfrom
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/167816
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit b65f821 with merge base 8f16199 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
release notes: onnx
| ).from_address(tensor.data_ptr()) | ||
|
|
||
| def tobytes(self) -> bytes: | ||
| # On big-endian machines, call the super's tobytes() which returns a little-endian result. |
There was a problem hiding this comment.
Does this work for types like bfloat16?
There was a problem hiding this comment.
@justinchuby I just tested with the following simple example, and it works for bfloat16.
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.tensor([5.0], dtype=torch.bfloat16))
def forward(self, x):
return x * self.weight
|
Could you fix the lint? |
mikaylagawarecki
added
the
triaged
@titaiwangms I didn't touch the file |
Try rebase maybe? |
Signed-off-by: Tung D. Le <tung@jp.ibm.com>
tungld
force-pushed
the
fix-onnx-export-big-endian
branch
from
7cf68a8 to
b65f821
Compare
|
@titaiwangms Thanks for the suggestion! I rebased it. Hope it works. |
|
@pytorchbot merge |
|
This PR needs to be approved by an authorized maintainer before merge. |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
6 participants
On big endian machines, constant values in the exported onnx model are still in big endian, they need to be converted to little endian to comply with the onnx specification.
This fixes that issue by calling super's methods of
ir.Tensorthat already handles endianness well.