add output for train loss in assertian err

2025-04-18 08:11:11 -07:00
parent 08f287b57f
commit a0670abc94
12 changed files with 25 additions and 25 deletions
--- a/tests/e2e/test_llama_pretrain.py
+++ b/tests/e2e/test_llama_pretrain.py
@@ -84,5 +84,5 @@ class TestPretrainLlama:
            temp_dir + "/runs",
            "train/train_loss",
            loss_threshold,
-            "Train Loss is too high",
+            "Train Loss (%s) is too high",
        )