Fix build crash in nvcr.io/nvidia/pytorch:24.06-py3 image (#2964)

InternLM · Dec 30, 2024 · b7488dd · b7488dd
1 parent af0fcf2
commit b7488dd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.cu b/src/turbomind/kernels/gemm/moe_utils_v2.cu
@@ -2,6 +2,7 @@
 
 #include <algorithm>
 #include <cstdio>
+#include <iostream>
 #include <limits>
 #include <numeric>
 #include <random>

diff --git a/src/turbomind/kernels/gemm/test/test_utils.cu b/src/turbomind/kernels/gemm/test/test_utils.cu
@@ -84,7 +84,7 @@ FastCompare(const T* src, const T* ref, int dims, int bsz, cudaStream_t stream,
         thrust::cuda::par.on(stream),
         zip_iter,
         zip_iter + count,
-        [=] __device__(auto tup) {
+        [=] __host__ __device__(thrust::tuple<float, float> tup) -> Tuple {
             float   s        = thrust::get<0>(tup);
             float   r        = thrust::get<1>(tup);
             float   abs_diff = fabsf(s - r);