add peak_memory info to training callback (#572)

This commit is contained in:
madroid 2024-03-14 11:17:10 +08:00 committed by GitHub
parent 376bb9cc44
commit d4e1de1d5b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -250,6 +250,7 @@ def train(
"iterations_per_second": it_sec,
"tokens_per_second": tokens_sec,
"trained_tokens": trained_tokens,
"peak_memory": peak_mem,
}
training_callback.on_train_loss_report(train_info)