提交 7aeb4f6c 编写于 作者: M Megvii Engine Team 提交者: Xu Xinran

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

GitOrigin-RevId: 85298084a337b4caa4aafb601259cf18ef3df557
上级 7a0c7ef4
......@@ -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.
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册