Skip to content

Commit

Permalink
[js/webgpu] Fix the crash issue in unsqueeze
Browse files Browse the repository at this point in the history
While allowing axes in unsqueeze to be scalar, its shape couldn't be
always accessed like a vector. This PR fixes issue microsoft#22031 so that the
original model could run well.
  • Loading branch information
Yang Gu committed Sep 29, 2024
1 parent 7880342 commit 280e3e0
Showing 1 changed file with 2 additions and 3 deletions.
5 changes: 2 additions & 3 deletions onnxruntime/core/providers/js/operators/unsqueeze.h
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,8 @@ class Unsqueeze final : public JsKernel, public UnsqueezeBase {
ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 0 ||
axes_tensor->Shape().NumDimensions() == 1,
"An axes tensor must be a scalar or a vector tensor.");
auto nDims = static_cast<size_t>(axes_tensor->Shape()[0]);
const auto* data = axes_tensor->Data<int64_t>();
axes.assign(data, data + nDims);
auto data_span = axes_tensor->template DataAsSpan<int64_t>();
axes.assign(data_span.begin(), data_span.end());
} else {
axes.assign(axes_.begin(), axes_.end());
}
Expand Down

0 comments on commit 280e3e0

Please sign in to comment.