Browse Source

fix(mge/optimizer): use static key to avoid mem leak

GitOrigin-RevId: 85298084a3
tags/v0.5.0
Megvii Engine Team Xu Xinran 5 years ago
parent
commit
7aeb4f6ca7
1 changed files with 3 additions and 3 deletions
  1. +3
    -3
      python_module/megengine/optimizer/optimizer.py

+ 3
- 3
python_module/megengine/optimizer/optimizer.py View File

@@ -225,13 +225,13 @@ class Optimizer(metaclass=ABCMeta):
param.grad.reset_zero()

def bcast_param(self):
key = 0
for group in self.param_groups:
for param in group["params"]:
bcast_param(
param,
"bcast_param_" + str(get_group_id()),
is_root=(get_rank() == 0),
param, "bcast_param_" + str(key), is_root=(get_rank() == 0),
)
key += 1

def state_dict(self) -> Dict:
r"""Export the optimizer state.


Loading…
Cancel
Save