From 7aeb4f6ca7f6b615fcf706647df20dc42386cfe7 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Mon, 15 Jun 2020 14:49:47 +0800 Subject: [PATCH] fix(mge/optimizer): use static key to avoid mem leak GitOrigin-RevId: 85298084a337b4caa4aafb601259cf18ef3df557 --- python_module/megengine/optimizer/optimizer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python_module/megengine/optimizer/optimizer.py b/python_module/megengine/optimizer/optimizer.py index 86c02c92b..b89eb1d1d 100644 --- a/python_module/megengine/optimizer/optimizer.py +++ b/python_module/megengine/optimizer/optimizer.py @@ -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. -- GitLab