Skip to content

[ONNX] Fix onnx constant folding #101329

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 6 commits into from
Closed

[ONNX] Fix onnx constant folding #101329

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cf607dc
fix onnx constant folding
May 13, 2023
d1d6602
remove debug print
May 15, 2023
72e8ec9
fix lint
May 15, 2023
0328d82
fix comment
May 16, 2023
0ecb534
comments
Jun 10, 2023
5746935
Update torch/csrc/jit/passes/onnx/constant_fold.cpp
BowenBao Jun 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions torch/csrc/jit/passes/onnx/constant_fold.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,6 +487,12 @@ c10::optional<at::Tensor> runTorchBackendForOnnx(
if (q > 1) {
return c10::nullopt;
}
// If the device of indices tensor is not the same with it of the input
// tensor, move it to the device of the input tensor
auto indices_val = node->input(1);
if (inputTensorValues[0].device() != indices.device()) {
indices = indices.to(inputTensorValues[0].device());
}
// If indices input for onnx::Gather has a value less than 0,
// It needs to be adjusted (+= dim value) for aten op
auto less_mask = at::lt(indices, 0);
Expand Down