Pass model's current output to the next step as recurrent input

[rushi-the-neural-arch]

I am working on a Vision problem where I am passing a 5D Tensor (for temporal information) as input to my model -
Batch_Size x Frames x C x H x W and I want to pass the model's recurrent state / intermediate level outputs continually to the next forward steps. That can be done easily by passing the "hidden" states like this in the training_step() and forward function-

training_step(self, batch, batch_idx, hiddens): forward(self, x, hiddens):

But the issue is for that, you need to enable truncated_bptt_steps. And as per my understanding keeping self.truncated_bptt_steps>0 passes only a partial input to the model, in my case it passes only 2 x 1 x 3 x 256 x 256 rather than 2 x 15 x 3 x 256 x 256

Sorry for the long info, but in summary I would like to pass the recurrent outputs of my models onto the next forward steps without truncating the input to my forward function, thanks in advance for any help!

Current boiler plate for the above operations -

class Trainer(pl.LightningModule):
    def __init__(self,):
        super(Trainer, self).__init__()

                self.truncated_bptt_steps = 1

def forward(self, x, hiddens):

    preds, *hiddens = self.model(x, *hiddens)
    return preds, *hiddens
        
def training_step(self, batch, batch_idx, hiddens):

    if batch_idx == 0:
        hiddens = [None] * 4
        print("-- For the First time --", hiddens)
    
    input, target = batch[0], batch[1]
        
    pred, *hiddens = self.forward(input, hiddens)
    loss = self.loss_func(pred, target)
    
    return {'loss' : loss.mean(), 'hiddens' : hiddens}

And if you do not pass self.truncated_bptt_steps=1 (>0) then it throws you the following error -
TypeError: training_step() missing 1 required positional argument: 'hiddens'

Is there any way to do this without using bptt?? Thanks!

References - How to add a hidden state in test_step when dealing with sequential data? #5785

Documentation