diff --git a/src/executor/graph_executor.cc b/src/executor/graph_executor.cc
index da1f13bce6c6..0b8a7f7d5de0 100644
--- a/src/executor/graph_executor.cc
+++ b/src/executor/graph_executor.cc
@@ -1022,8 +1022,8 @@ void GraphExecutor::InitDataEntryMemory(std::vector<NDArray>* shared_pool) {
     for (uint32_t i = 0; i < idx[nid].source->num_outputs(); ++i) {
       auto eid = idx.entry_id(nid, i);
       data_context[eid] = vctx[nid];
-      CHECK_NE(vstorage_type[nid], kUndefinedStorage);
-      data_storage_type[eid] = (NDArrayStorageType) vstorage_type[nid];
+      CHECK_NE(vstorage_type[eid], kUndefinedStorage);
+      data_storage_type[eid] = (NDArrayStorageType) vstorage_type[eid];
     }
   }