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

pull/1423/head
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

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