review feedback

microsoft · Jan 9, 2024 · 6e3182e · 6e3182e
1 parent 10db421
commit 6e3182e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/js/web/lib/wasm/jsep/webgpu/ops/where.ts b/js/web/lib/wasm/jsep/webgpu/ops/where.ts
@@ -74,9 +74,8 @@ const createWhereOpProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
   const outputDataType = inputs[1].dataType;
 
   const isBroadcast = !(ShapeUtil.areEqual(dimsA, dimsB) && ShapeUtil.areEqual(dimsB, dimsC));
-  let outputShape = dimsC;
-  let outputSize = ShapeUtil.size(dimsC);
-  const vecSize = Math.ceil(outputSize / 4);
+  let outputShape = dimsA;
+  let outputSize = ShapeUtil.size(dimsA);
   // TODO: deal with zero-sized tensors (eg. dims=[1,0])
 
   if (isBroadcast) {
@@ -88,6 +87,8 @@ const createWhereOpProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
     outputSize = ShapeUtil.size(outputShape);
   }
 
+  const vecSize = Math.ceil(outputSize / 4);
+
   return {
     name: 'Where',
     shaderCache: {inputDependencies: ['rank', 'rank', 'rank']},