Fix NDArray ToDLPack Bug (#13698)

apache · Dec 21, 2018 · f211588 · f211588
1 parent 63fe849
commit f211588
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
@@ -330,11 +330,10 @@ struct NDArrayDLManager {
 };
 
 DLManagedTensor* NDArray::ToDLPack() const {
+  CHECK(!is_none()) << "NDArray is not initialized";
   NDArrayDLManager* dlmanager(new NDArrayDLManager);
   dlmanager->handle = *this;
-  if (!is_none()) {
-    dlmanager->tensor.dl_tensor = data().dltensor();
-  }
+  dlmanager->tensor.dl_tensor = dlmanager->handle.data().dltensor();
   dlmanager->tensor.manager_ctx = dlmanager;
   dlmanager->tensor.deleter = [](DLManagedTensor* dlmanager){
     delete static_cast<NDArrayDLManager*>(dlmanager->manager_ctx);