fix f16 for attention, enable slice and flatten for more types (micro…

…soft#19262)
siweic0 · Jan 29, 2024 · f01392c · f01392c
1 parent 1161753
commit f01392c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/web/lib/wasm/jsep/webgpu/ops/attention.ts b/web/lib/wasm/jsep/webgpu/ops/attention.ts
@@ -297,7 +297,7 @@ export const computeInPlaceSoftmax = (context: ComputeContext, input: TensorView
 
     if (sum == 0) {
       for (var i: u32 = 0; i < uniforms.elements_per_wg && i + localOffset < uniforms.d_comp; i++) {
-        x[offset + i] = ${fillVector('f32', components, 'uniforms.d_inv')};
+        x[offset + i] = ${fillVector(elemValueType, components, 'uniforms.d_inv')};
       }
     } else {
       for (var i: u32 = 0; i < uniforms.elements_per_wg && i + localOffset < uniforms.d_comp; i++) {