Corrected the displayed loss when gradient_accumulation_steps > 1

8d1d1ffd · Matthew Carrigan · abb7d1ff · 8d1d1ffd
Commit 8d1d1ffd authored Mar 25, 2019 by Matthew Carrigan
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

examples/lm_finetuning/finetune_on_pregenerated.py examples/lm_finetuning/finetune_on_pregenerated.py +1 -1

No files found.
--- a/examples/lm_finetuning/finetune_on_pregenerated.py
+++ b/examples/lm_finetuning/finetune_on_pregenerated.py
@@ -309,7 +309,7 @@ def main():
                nb_tr_examples += input_ids.size(0)
                nb_tr_steps += 1
                pbar.update(1)
-                mean_loss = tr_loss / nb_tr_steps
+                mean_loss = tr_loss * args.gradient_accumulation_steps / nb_tr_steps
                pbar.set_postfix_str(f"Loss: {mean_loss:.5f}")
                if (step + 1) % args.gradient_accumulation_steps == 0:
                    if args.fp16: