fix block.export (#17970)

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)