Dataloader reinsertion for recursive predictionsrs

[leonardcaquot94]

I'm working with recursive predictions of variable lengths (e.g., sequence predictions) and exploring two approaches for handling backpropagation:

1. In the first approach, at each training_step, I predict the entire sequence recursively, accumulate the losses at each iteration, and return a single loss at the end of the loop. I then perform one update for the entire recursive prediction.
2. In the second approach, at each training_step, I predict only one step of the recursive process, allowing each step to trigger an update.

To implement this second approach, I need a way to reinsert partially completed predictions back into the dataloader to continue processing them in future batches. Since the sequences I’m predicting have a maximum length, there’s no risk of endlessly reinserting them—they will either reach their max size or the end-of-sequence token will be predicted.

Is there a straightforward way to implement this using PyTorch Lightning?