Use new api to get loss scale when finetuning.

a13cbe1e · Jared Casper · 7381754e · a13cbe1e
Commit a13cbe1e authored Jan 05, 2021 by Jared Casper
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

tasks/finetune_utils.py tasks/finetune_utils.py +2 -1

No files found.
--- a/tasks/finetune_utils.py
+++ b/tasks/finetune_utils.py
@@ -186,7 +186,8 @@ def _train(model, optimizer, lr_scheduler, forward_step,
            # Logging.
            report_memory_flag = training_log(losses_dict, losses_dict_sum,
                                              optimizer.param_groups[0]['lr'],
-                                              iteration, optimizer.loss_scale,
+                                              iteration,
+                                              optimizer.get_loss_scale().item(),
                                              report_memory_flag, skipped_iter)

            # Autoresume