slight changes

This commit is contained in:
Dan Saunders
2025-03-05 16:25:49 +00:00
parent 035e7a2f4c
commit e735f4270b
3 changed files with 5 additions and 4 deletions

View File

@@ -113,7 +113,7 @@ class TestTelemetryCallback:
callback.on_train_begin(training_args, trainer_state, trainer_control)
mock_telemetry_manager.send_event.assert_called_once_with(
event_type="train-started"
event_type="train-start"
)
def test_on_train_end(
@@ -130,7 +130,7 @@ class TestTelemetryCallback:
mock_telemetry_manager.send_event.assert_called_once()
call_args = mock_telemetry_manager.send_event.call_args[1]
assert call_args["event_type"] == "train-ended"
assert call_args["event_type"] == "train-end"
assert "loss" in call_args["properties"]
assert call_args["properties"]["loss"] == 2.5
assert "learning_rate" in call_args["properties"]