Update trainer.py

FlagAI-Open · Apr 28, 2024 · a6128ae · a6128ae
1 parent 5ecde7e
commit a6128ae
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flagai/trainer.py b/flagai/trainer.py
@@ -701,7 +701,7 @@ def train_step_pytorch(self,
                     optimizer.zero_grad()
                 else:
                     optimizer.step()
-                    # optimizer.zero_grad()
+                    optimizer.zero_grad()
                 self.accumulate_count = 0
             else:
                 self.accumulate_count += 1
@@ -775,7 +775,7 @@ def train_step_pytorchDDP(self,
                         optimizer.zero_grad()
                     else:
                         optimizer.step()
-                        # model.zero_grad()
+                        model.zero_grad()
 
                     self.accumulate_count = 0
                 else: