fix block.export

python/mxnet/gluon/block.py

@@ -1348,8 +1348,17 @@ 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.endswith('running_mean') or name.endswith('running_var') \
+                        or name.endswith('moving_mean') or name.endswith('moving_var'):
+                        if name not in aux_names:
+                            warnings.warn('Parameter "{name}" is not found in '
+                                          'the graph. '
+                                          .format(name=name), stacklevel=3)
+                    else:
+                        warnings.warn('Parameter "{name}" is not found in '
+                                      'the graph. '
+                                      .format(name=name), stacklevel=3)
+                    arg_dict['aux:%s'%name] = param._reduce()
         save_fn = _mx_npx.save if is_np_array() else ndarray.save
         params_filename = '%s-%04d.params'%(path, epoch)
         save_fn(params_filename, arg_dict)