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

Commit

Permalink
fix block.export (#17970) (#19075)
Browse files Browse the repository at this point in the history
This PR cherry-picks commit 5122d32 into the v1.x branch. This is to enable the export of models where dangling layers are optimized out during symbol export. For more information, see here and here.
  • Loading branch information
James Mracek committed Sep 3, 2020
1 parent b5e9c99 commit 748eebd
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions python/mxnet/gluon/block.py
Expand Up @@ -1238,8 +1238,11 @@ def export(self, path, epoch=0, remove_amp_cast=True):
if name in arg_names:
arg_dict['arg:{}'.format(name)] = param._reduce()
else:
assert name in aux_names
arg_dict['aux:{}'.format(name)] = param._reduce()
if name not in aux_names:
warnings.warn('Parameter "{name}" is not found in the graph. '
.format(name=name), stacklevel=3)
else:
arg_dict['aux:%s'%name] = param._reduce()
save_fn = _mx_npx.save if is_np_array() else ndarray.save
save_fn('%s-%04d.params'%(path, epoch), arg_dict)

Expand Down

0 comments on commit 748eebd

Please sign in to comment.