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_reward_model_smollm2.py
+++ b/tests/e2e/test_reward_model_smollm2.py
@@ -73,6 +73,6 @@ class TestRewardModelLoraSmolLM2(unittest.TestCase):

        train(cfg=cfg, dataset_meta=dataset_meta)
        check_tensorboard(
-            temp_dir + "/runs", "train/train_loss", 2.5, "Train Loss is too high"
+            temp_dir + "/runs", "train/train_loss", 2.5, "Train loss (%s) is too high"
        )
        check_model_output_exists(temp_dir, cfg)