mirror of https://github.com/hpcaitech/ColossalAI
[hotfix] zero optim prevents calling inner optim.zero_grad (#1422)
parent
504419d261
commit
6df3e19be9
|
@ -139,6 +139,7 @@ class ZeroOptimizer(ColossalaiOptimizer):
|
|||
self._update_params_ptr()
|
||||
ret = self.optim.step(*args, **kwargs)
|
||||
self._register_states()
|
||||
self.zero_grad()
|
||||
self._update_fp16_params()
|
||||
return ret
|
||||
|
||||
|
|
Loading…
Reference in New Issue