Dont divide by vec.size

d9867f64 · Paul · 84d0f5c9 · d9867f64
Commit d9867f64 authored Jul 04, 2022 by Paul
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/targets/gpu/jit/layernorm.cpp src/targets/gpu/jit/layernorm.cpp +1 -1

No files found.
--- a/src/targets/gpu/jit/layernorm.cpp
+++ b/src/targets/gpu/jit/layernorm.cpp
@@ -61,7 +61,7 @@ struct layernorm_compiler : compiler<layernorm_compiler>
        }
        auto preloads   = preload::broadcasts(axis, inputs);
        auto relements  = inputs[0].lens()[axis] / vec.size;
-        auto nelements  = (inputs.back().elements() / inputs[0].lens()[axis]) / vec.size;
+        auto nelements  = (inputs.back().elements() / inputs[0].lens()[axis]);
        auto block_size = compute_block_size(relements, 256);
        hip_compile_options options;
        options.set_launch_params(