adding missing field (facebookresearch#162)

"total_loss" was not saved
jianan-gu · Feb 26, 2021 · d25bbc2 · d25bbc2
1 parent 1302c71
commit d25bbc2
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/dlrm_s_pytorch.py b/dlrm_s_pytorch.py
@@ -1668,6 +1668,7 @@ def run():
                             model_metrics_dict["epoch"] = k
                             model_metrics_dict["iter"] = j + 1
                             model_metrics_dict["train_loss"] = train_loss
+                            model_metrics_dict["total_loss"] = total_loss
                             model_metrics_dict[
                                 "opt_state_dict"
                             ] = optimizer.state_dict()