Custom train loop to perform partial batch updates

[leonardcaquot94]

Hi there,

I'm working on a recursive data processing task in PyTorch Lightning, where the batch size decreases with each iteration as predictions are made until no more are needed. I want to address two key issues:

1. Instead of waiting until all iterations on a single batch are complete before I update the model with a total loss, I want to backpropagate after each iteration to make more frequent updates (a bit like in reinforcement learning problems, update at each step instead of each episode).
2. To avoid inefficiencies from shrinking batch sizes, I’d like to refill the batch by replacing completed elements with new data during each iteration, ensuring the batch size remains constant.

I think I need to change how training_step and validation_step work. Instead of iterating over a single batch until all recursive predictions are made, these methods should perform just one step at a time. After each step, completed predictions should be removed from the batch, while the remaining elements stay for further processing. This way, I can keep the batch size more consistent and improve efficiency.

Any idea that can help ?