Fix CopyVectorToNDArray in src/c_api_common.h (dmlc#3597)

* fix CopyVectorToNDArray * Fix lint Co-authored-by: Jinjing Zhou <[email protected]>
anirband · Dec 19, 2021 · 25538ba · 25538ba
1 parent a62f2c1
commit 25538ba
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/c_api_common.h b/src/c_api_common.h
@@ -31,16 +31,18 @@ dgl::runtime::PackedFunc ConvertNDArrayVectorToPackedFunc(
     const std::vector<dgl::runtime::NDArray>& vec);
 
 /*!
- * \brief Copy a vector to an int64_t NDArray.
+ * \brief Copy a vector to an NDArray.
  *
- * The element type of the vector must be convertible to int64_t.
+ * The data type of the NDArray will be IdType, which must be an integer type.
+ * The element type (DType) of the vector must be convertible to IdType.
  */
 template<typename IdType, typename DType>
 dgl::runtime::NDArray CopyVectorToNDArray(
     const std::vector<DType>& vec) {
   using dgl::runtime::NDArray;
   const int64_t len = vec.size();
-  NDArray a = NDArray::Empty({len}, DLDataType{kDLInt, sizeof(IdType), 1}, DLContext{kDLCPU, 0});
+  NDArray a = NDArray::Empty({len}, DLDataType{kDLInt, sizeof(IdType) * 8, 1},
+                             DLContext{kDLCPU, 0});
   std::copy(vec.begin(), vec.end(), static_cast<IdType*>(a->data));
   return a;
 }