Fix on logging embedding norms data. Save steps now matches logging steps.

joaomonteirof · joaomonteirof · commit 1496a642894f · 2023-03-22T16:12:27.000Z
diff --git a/src/hf_trainer.py b/src/hf_trainer.py
@@ -44,7 +44,7 @@ def compute_metrics(eval_pred: PredictionOutput) -> Dict[str, float]:
         "R@1": recall_at_1.item(),
         "R@5": recall_at_5.item(),
         "MRR": mean_reciprocal_rank.item(),
-        "embedding_norms": embedding_norms,
+        "embedding_norms": [norm for norm in embedding_norms],
         "min_embedding_norm": embedding_norms.min().item(),
     }
 
@@ -263,6 +263,7 @@ def get_trainer(
         logging_strategy="steps",
         logging_steps=log_every,
         save_strategy="steps",
+        save_steps=log_every,
         evaluation_strategy="steps",
         report_to="wandb",
     )
