From 6df3e19be9240ca954c56215bb48ff3667825122 Mon Sep 17 00:00:00 2001
From: ver217 <lhx0217@gmail.com>
Date: Tue, 9 Aug 2022 16:08:12 +0800
Subject: [PATCH] [hotfix] zero optim prevents calling inner optim.zero_grad
 (#1422)

---
 colossalai/zero/zero_optimizer.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/colossalai/zero/zero_optimizer.py b/colossalai/zero/zero_optimizer.py
index e5a2f9f90..bc23123b0 100644
--- a/colossalai/zero/zero_optimizer.py
+++ b/colossalai/zero/zero_optimizer.py
@@ -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