LSTM training with hiddens in training_step()
I am trying an LSTM network where I require in each training step the hidden value to be passed over to the next training step until the end of epoch. At the start of new epoch we reset by initializing the hidden to zeros. Following is how i do this,
def training_step(self, batch, batch_nb, hiddens):
        ...
        if batch_nb == 0: #init only in new epochs
            hiddens = self.model.init_hidden(self.hparams.batch_size)
        out, hiddens = self.forward(posts, hiddens)
        ...
        loss = self.criterion(out, targets)
        return {'loss': loss, 'hiddens': hiddens}

As per the docs suggest, I've added the truncated_bptt_steps argument to Trainer, but it does not work as expected. It gives the following exception,
AssertionError: Batch time dimension length is ambiguous
I've tried setting truncated_bptt_steps=2 to different values but what ever value i set it results in the same exception. How does one train LSTMs like this with PTL? What am i doing wrong here?
PTL version 1.1.2