Merge pull request #706 from mv1388/add_optimizer_zerograd_to_TestIMD…

…BBERT_AMP Add optmizer zero grad to the AMP BERT GPU test
mv1388 · Jul 16, 2022 · d904e34 · d904e34
2 parents 09a687b + e3300df
commit d904e34
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests_gpu/test_single_gpu/test_amp/test_imdb_bert.py b/tests_gpu/test_single_gpu/test_amp/test_imdb_bert.py
@@ -109,6 +109,8 @@ def train_eval_core_pytorch(self, train_data, test_data, num_epochs):
                 scaler.step(optimizer)
                 scaler.update()
 
+                optimizer.zero_grad()
+
             # Imitate what happens in auto_execute_end_of_epoch() in TrainLoop
             for _ in train_loader:
                 pass