From 434f0fa53614b297610d285f03e0215cbfb5de25 Mon Sep 17 00:00:00 2001 From: Yang Gu Date: Mon, 30 Sep 2024 17:28:16 +0800 Subject: [PATCH] [js/webgpu] Fix the crash issue in unsqueeze (#22264) 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. --- onnxruntime/core/providers/js/operators/unsqueeze.h | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/js/operators/unsqueeze.h b/onnxruntime/core/providers/js/operators/unsqueeze.h index f15a3008895aa..dd5f7e0525669 100644 --- a/onnxruntime/core/providers/js/operators/unsqueeze.h +++ b/onnxruntime/core/providers/js/operators/unsqueeze.h @@ -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(axes_tensor->Shape()[0]); - const auto* data = axes_tensor->Data(); - axes.assign(data, data + nDims); + auto data_span = axes_tensor->template DataAsSpan(); + axes.assign(data_span.begin(), data_span.end()); } else { axes.assign(axes_.begin(), axes_.end()); }