fix for ch11 (apache#15244)

haojin2 · Jul 31, 2019 · 3252e9b · 3252e9b
1 parent e92d044
commit 3252e9b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/python/mxnet/gluon/parameter.py b/python/mxnet/gluon/parameter.py
@@ -369,7 +369,7 @@ def _reduce(self):
         ctx = context.cpu()
         if self._stype == 'default':
             block = self.list_data()
-            data = ndarray.add_n(*(w.copyto(ctx) for w in block)) / len(block)
+            data = ndarray.add_n(*(w.copyto(ctx).as_nd_ndarray() for w in block)) / len(block)
         else:
             # fetch all rows for 'row_sparse' param
             all_row_ids = ndarray.arange(0, self.shape[0], dtype='int64', ctx=ctx)

diff --git a/python/mxnet/numpy_extension/__init__.py b/python/mxnet/numpy_extension/__init__.py
@@ -28,5 +28,6 @@
 from ..util import use_np_shape, np_shape, is_np_shape
 from ..util import use_np_array, np_array, is_np_array
 from ..util import set_np, use_np, reset_np
+from ..ndarray import waitall
 
 __all__ = []