Use a metric from a callback

[Turakar]

Hello,

I want to log throughput of my model in tokens per second using a metric. To this end, I had the following idea:

1. Register a callback.
2. This callback stores the start time of each batch in on_train_batch_start().
3. In on_train_batch_end(), we can compute the throughput for the local process.
4. Then, log the throughput using a MeanMetric, so that torchmetrics can take care of distributed communication.

However, what is unclear to me, is how to log the metric and how to ensure that it is properly reset. I want to log a value every Trainer.log_every_n_steps steps.

Thank you for your help.