fix test

Signed-off-by: Ananth Subramaniam <[email protected]>

Author: ananthsub

tests/unit_tests/data/datasets/test_sft.py

@@ -278,7 +278,7 @@ def test_collate_fn(self, ensure_test_data):
                 "answer_ids": np.array([104, 105]),
                 "metadata": {"id": "ex1"},
                 "seq_boundaries": (0, 3),
-                "mask": np.array([0, 0, 0, 1, 1]),
+                "loss_mask": np.array([0, 0, 0, 1, 1]),  # Changed from "mask" to "loss_mask"
                 "metadata": {},
                 "token_count": 5,
             },
@@ -290,7 +290,7 @@ def test_collate_fn(self, ensure_test_data):
                 "answer_ids": np.array([203, 204]),
                 "metadata": {"id": "ex2"},
                 "seq_boundaries": (0, 2),
-                "mask": np.array([0, 0, 1, 1]),
+                "loss_mask": np.array([0, 0, 1, 1]),  # Changed from "mask" to "loss_mask"
                 "metadata": {},
                 "token_count": 4,
             },