Fix the input_output_stride for 4-dim input.

microsoft · Jan 21, 2025 · a4d8482 · a4d8482
1 parent 1009fc9
commit a4d8482
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc
@@ -95,7 +95,7 @@ Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& con
       input_shape.NumDimensions() == 3
           ? std::vector<uint32_t>({batch_stride, hidden_size, head_size, 1})
           : (input_shape.NumDimensions() == 4
-                 ? std::vector<uint32_t>({batch_stride, head_size, sequence_length * head_size, 1})
+                 ? std::vector<uint32_t>({batch_stride, sequence_length * head_size, head_size, 1})
                  : std::vector<uint32_t>({}));
 
   program