fix a typo that fails the compiling when ElementScale is not the same…

… as MmaType (#1977)
NVIDIA · Dec 10, 2024 · 2b6cfd3 · 2b6cfd3
1 parent 4c42f73
commit 2b6cfd3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/include/cutlass/detail/collective/mixed_input_utils.hpp b/include/cutlass/detail/collective/mixed_input_utils.hpp
@@ -836,7 +836,7 @@ struct MixedInputUtils {
         }
       }
       else {
-        auto stage = make_tensor_like<ElementScale>(src_vm);
+        auto stage = make_tensor_like<ElementScale>(src_vm(_, 0));
         CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < size<1>(dst_vm); ++i) {
           LayoutAwareConvert(src_vm(_, i), stage);
@@ -868,7 +868,7 @@ struct MixedInputUtils {
         }
       }
       else {
-        auto stage = make_tensor_like<ElementScale>(src_vm);
+        auto stage = make_tensor_like<ElementScale>(src_vm(_, 0));
         CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < size<1>(dst_vm); ++i) {
           LayoutAwareConvert(src_vm(_, i), stage);