Save train predictions at the end of batch with Trainer

I am using Trainer to train my model.
I am looking for a way to store all predictions per batch and aggregate them at the end of the epoch in order to calculate train epoch accuracy / loss / other metrics (it maybe less accurate, but save computational resources compare to evaluating train dataset at the end of the epoch).
Is there a callback or other suggested solution? (…on_train_batch_end?)


You can use callback with trainer (Trainer)

Also, you have some many options with callback(Callbacks)


I am familiar with Callbacks option. I am looking for something like “on_batch_end”. I could not find something similar within the TrainerCallback build in functions. I want to collect all train predictions per batch.
Any idea?

There callback has on_epoch_end and on_step_end .

I think what you say is ‘on_epoch_end’ . link

Sorry, I meant on_batch_end.
Can you help?

I’m sorry, i also confuse with on_step_end and on_sub_step_end.

I think we need to build new custom callback for check each batch.

ok. Thanks.
If you build such callback, I will be glad if you can post it here :pray:t4: