fix lint

dmlc · Sep 19, 2023 · d0d5032 · d0d5032
1 parent 5a7558a
commit d0d5032
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/dgl_sparse/src/matrix_ops_impl.h b/dgl_sparse/src/matrix_ops_impl.h
@@ -10,6 +10,7 @@
 #include <sparse/sparse_matrix.h>
 
 #include <tuple>
+#include <vector>
 
 #include "./utils.h"
 
@@ -73,19 +74,18 @@ CompactImpl(
       nnz.push_back(i);
     }
   }
-  long int nnz_size = nnz.size();
   if (dim == 0) {
     auto ret = SparseMatrix::FromCSR(
         VectorToTorchTensor(ret_ptr), VectorToTorchTensor(ret_idx),
         VectorToTorchTensor(ret_val),
-        std::vector<int64_t>{nnz_size, csr->num_cols});
+        std::vector<int64_t>{static_cast<int64_t>(nnz.size()), csr->num_cols});
     auto ret_idx = torch::optional<torch::Tensor>(VectorToTorchTensor(nnz));
     return {ret, ret_idx};
   } else {
     auto ret = SparseMatrix::FromCSC(
         VectorToTorchTensor(ret_ptr), VectorToTorchTensor(ret_idx),
         VectorToTorchTensor(ret_val),
-        std::vector<int64_t>{csr->num_cols, nnz_size});
+        std::vector<int64_t>{csr->num_cols, static_cast<int64_t>(nnz.size())});
     auto ret_idx = torch::optional<torch::Tensor>(VectorToTorchTensor(nnz));
     return {ret, ret_idx};
   }

diff --git a/dgl_sparse/src/utils.h b/dgl_sparse/src/utils.h
@@ -15,6 +15,8 @@
 #include <torch/custom_class.h>
 #include <torch/script.h>
 
+#include <vector>
+
 namespace dgl {
 namespace sparse {