Skip to content

Commit

Permalink
[js/webgpu] Fix the crash issue in unsqueeze (#22264)
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 #22031 so that the
original model could run well.
  • Loading branch information
Yang Gu authored Sep 30, 2024
1 parent 1bda91f commit 434f0fa
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 434f0fa

Please sign in to comment.