From fc52b29b2e692b2cf04ab3309c4cef37fc0cee08 Mon Sep 17 00:00:00 2001
From: Hao Lu <hlu@fb.com>
Date: Thu, 9 Apr 2020 01:33:10 -0700
Subject: [PATCH] [NDArray] Set shape_ in NDArray::FromDLPack

---
 src/runtime/ndarray.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/runtime/ndarray.cc b/src/runtime/ndarray.cc
index 99594ee50f93..ac12472a903e 100644
--- a/src/runtime/ndarray.cc
+++ b/src/runtime/ndarray.cc
@@ -208,6 +208,10 @@ NDArray NDArray::FromDLPack(DLManagedTensor* tensor) {
   // fill up content.
   data->manager_ctx = tensor;
   data->dl_tensor = tensor->dl_tensor;
+  // update shape_
+  data->shape_.resize(data->dl_tensor.ndim);
+  data->shape_.assign(data->dl_tensor.shape, data->dl_tensor.shape + data->dl_tensor.ndim);
+  data->dl_tensor.shape = data->shape_.data();
   return NDArray(GetObjectPtr<Object>(data));
 }