Update softmax

77ebc94b · Paul · e3e8dc32 · 77ebc94b
Commit 77ebc94b authored Nov 17, 2022 by Paul
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

src/targets/gpu/kernels/include/migraphx/kernels/softmax.hpp src/targets/gpu/kernels/include/migraphx/kernels/softmax.hpp +2 -1

No files found.
--- a/src/targets/gpu/kernels/include/migraphx/kernels/softmax.hpp
+++ b/src/targets/gpu/kernels/include/migraphx/kernels/softmax.hpp
@@ -30,9 +30,10 @@
 namespace migraphx {
 template <index_int Axis, class Input, class Output>
-__device__ void softmax(Input input, Output output)
+__device__ void softmax(Input input1, Output output)
 {
    reduce::block::run<reduce::with_axis<Input, Axis>>([&](auto, auto r) {
+        auto input = r.inner(op::id{})(input1);
 #ifdef MIGRAPHX_USE_FAST_SOFTMAX
        const auto c = vec_at(r.slice(input)[0], 0);
 #else