Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Enable Numpy support for Gluon Block optimize_for #19455

Merged
merged 4 commits into from
Nov 2, 2020
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions python/mxnet/gluon/block.py
Original file line number Diff line number Diff line change
Expand Up @@ -1038,6 +1038,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

Expand Down