Browse Source

fix(mge/parampack): fix param pack when no param left

GitOrigin-RevId: b88b876064
release-1.1
Megvii Engine Team 4 years ago
parent
commit
a1efd4d56d
1 changed files with 2 additions and 0 deletions
  1. +2
    -0
      imperative/python/megengine/distributed/helper.py

+ 2
- 0
imperative/python/megengine/distributed/helper.py View File

@@ -211,6 +211,8 @@ class AllreduceCallback:
self._grad_origin_device = dict()

def _pack(self, dtype):
if len(self._packing_list[dtype]) == 0:
return
grad_list = [self._gradients_dict[p] for p in self._packing_list[dtype]]
shapes = [p.shape for p in self._packing_list[dtype]]
reduced_grads = pack_allreduce_split(


Loading…
Cancel
Save