Fix GetMKLDNNData for delay alloc (#16618)

* Fix GetMKLDNNData for delay alloc * Run CI * Run CI * Run CI * Run CI * Run CI Change-Id: I7ac2796e0ee8439c92fd2bd7a70a23a359b76b12
apache · Oct 28, 2019 · 11dff51 · 11dff51
1 parent ffc5392
commit 11dff51
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
@@ -647,6 +647,7 @@ const mkldnn::memory *NDArray::GetMKLDNNData() const {
     // If this is a view, we can't create a MKLDNN memory for the chunk
     // because we don't have the complete data type and shape information for
     // the chunk.
+    CheckAndAlloc();
     void *off_addr = static_cast<char *>(ptr_->shandle.dptr) + byte_offset_;
     // Create the primitive desc for the new mkldnn memory.
     mkldnn::memory::dims dims(shape().ndim());
@@ -665,6 +666,7 @@ const mkldnn::memory *NDArray::GetMKLDNNData() const {
   } else {
     // If this isn't a view, we can create a MKLDNN memory and store it in the
     // chunk.
+    CheckAndAlloc();
     ptr_->SetMKLMem(shape_, dtype_);
     MKLDNNStream::Get()->RegisterMem(ptr_->mkl_mem_->GetMem());
     return ptr_->mkl_mem_->GetRaw();