github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-16 08:27:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] zero optim prevents calling inner optim.zero_grad (#1422)

Browse Source
This commit is contained in:
ver217 2022-08-09 16:08:12 +08:00 committed by GitHub
parent 504419d261
commit 6df3e19be9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/zero/zero_optimizer.py
View File

@ -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