Divide by vec.size

84d0f5c9 · Paul · 047162bb · 84d0f5c9
Commit 84d0f5c9 authored Jul 02, 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];
+        auto nelements  = (inputs.back().elements() / inputs[0].lens()[axis]) / vec.size;
        auto block_size = compute_block_size(relements, 256);
        hip_compile_options options;
        options.set_launch_params(