From 748eebd7475b3052d5cc966765c6913cb181ad4f Mon Sep 17 00:00:00 2001 From: James Mracek Date: Thu, 3 Sep 2020 09:21:50 -0700 Subject: [PATCH] fix block.export (#17970) (#19075) 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. --- python/mxnet/gluon/block.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py index 9772e2394486..41ef2cb15d89 100644 --- a/python/mxnet/gluon/block.py +++ b/python/mxnet/gluon/block.py @@ -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)