initial commit

apache · Oct 30, 2020 · 16038a1 · 16038a1
1 parent 3b8b0c2
commit 16038a1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py
@@ -966,6 +966,10 @@ def _build_cache(self, *args):
             # Partition the graph.
             out = out.optimize_for(self._backend, arg_dict, aux_dict, ctx, **self._backend_opts)
 
+            # convert to numpy symbol if needed
+            if _mx_npx.is_np_array():
+                out = out.as_np_ndarray()
+
             #update cached graph with partitioned graph
             self._cached_graph = data, out