mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-25 01:41:19 +08:00
add peak_memory info to training callback (#572)
This commit is contained in:
parent
376bb9cc44
commit
d4e1de1d5b
@ -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)
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user