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: