Improve vectorization

src/targets/gpu/jit/fused_reduce.cpp

@@ -124,7 +124,7 @@ struct fused_reduce_compiler : compiler<fused_reduce_compiler>
         if(algo == "block")
         {
             // Vectorize if the axis is a reduction axis
-            if(options.virtual_inputs.back().lens()[faxis] == 1)
+            if(output_shape.lens()[faxis] == 1)
                 vec = vectorize::elements(ctx, faxis, options.virtual_inputs);
             auto relements  = reduced_shape.elements() / vec.size;
             auto block_size = compute_block_size(relements, 256);