saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-06-09 00:30:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/webgpu] Fix the crash issue in unsqueeze (#22264)

Browse source 
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.
This commit is contained in:
Yang Gu 2024-09-30 17:28:16 +08:00 committed by GitHub
parent 1bda91fc57
commit 434f0fa536
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
onnxruntime/core/providers/js/operators/unsqueeze.h
View file

@ -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());
}