add check in cuda file

dmlc · Jun 13, 2024 · 86f8846 · 86f8846
1 parent 14ac2ef
commit 86f8846
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/array/cuda/coo2csr.cu b/src/array/cuda/coo2csr.cu
@@ -40,6 +40,7 @@ CSRMatrix COOToCSR<kDGLCUDA, int32_t>(COOMatrix coo) {
   }
 
   const int64_t nnz = coo.row->shape[0];
+  CHECK_NO_OVERFLOW(coo.row->dtype, nnz);
   // TODO(minjie): Many of our current implementation assumes that CSR must have
   //   a data array. This is a temporary workaround. Remove this after:
   //   - The old immutable graph implementation is deprecated.