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)