set context for newly created aux handle

apache · Mar 12, 2019 · 2248647 · 2248647
1 parent 1f4626b
commit 2248647
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/include/mxnet/ndarray.h b/include/mxnet/ndarray.h
@@ -1051,6 +1051,8 @@ class NDArray {
         << "storage type cannot be kDefaultStorage in CheckAndAllocAuxData";
       if (aux_handles.size() <= i) {
         aux_handles.resize(i + 1);
+        // set context for the newly created handle.
+        aux_handles[i].ctx = ctx;
       }
       size_t aux_bytes = shape.Size() * mshadow::mshadow_sizeof(aux_types[i]);
       if (aux_handles[i].size < aux_bytes) {