diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py index bed6679be2e6..d61dbaddbc7b 100644 --- a/python/mxnet/gluon/block.py +++ b/python/mxnet/gluon/block.py @@ -1195,12 +1195,16 @@ def export(self, path, epoch=0, remove_amp_cast=True): arg_names = set(sym.list_arguments()) aux_names = set(sym.list_auxiliary_states()) arg_dict = {} - for name, param in self.collect_params().items(): - if name in arg_names: - arg_dict['arg:%s'%name] = param._reduce() - else: - assert name in aux_names - arg_dict['aux:%s'%name] = param._reduce() + for is_arg, name, param in self._cached_op_args: + if not is_arg: + if name in arg_names: + arg_dict['arg:{}'.format(name)] = param._reduce() + else: + 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)