Dont divide by vec.size

e01224f5 · Paul · 0a5e9b99 · e01224f5
Commit e01224f5 authored Jul 04, 2022 by Paul
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

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

No files found.
--- a/src/targets/gpu/jit/softmax.cpp
+++ b/src/targets/gpu/jit/softmax.cpp
@@ -80,7 +80,7 @@ struct softmax_compiler : compiler<softmax_compiler>
            vec = vectorize::elements(faxis, 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(