Fit Error: raised training_step() takes 3 positional arguments but 4 were given when I use truncated_bptt_steps
🐛 Bug
Everything works fine when I had truncated_bptt_steps as None, but when I set it to 5. The error mentioned in the title is thrown (see below for Traceback detail):
Traceback (most recent call last):
  File "Trainer.py", line 103, in <module>
    trainer.fit(gen)
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 695, in fit
    self.single_gpu_train(model)
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 441, in single_gpu_train
    self.run_pretrain_routine(model)
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 829, in run_pretrain_routine
    self.train()
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 332, in train
    self.run_training_epoch()
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 386, in run_training_epoch
    output = self.run_training_batch(batch, batch_idx)
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 506, in run_training_batch
    loss = optimizer_closure()
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 475, in optimizer_closure
    split_batch, batch_idx, opt_idx, self.hiddens)
  File "xxxx/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 593, in training_forward
    output = self.model.training_step(*args)
TypeError: training_step() takes 3 positional arguments but 4 were given

Code sample
My DataSet class
class DataSet(data.Dataset):
    def __init__(self, data, token_id2idx):
        '''
        Data and label has to have the same time step size for truncated batch split to work
        'label is the data offset by 1 (because dont need to predict the first CLS token)
        :param data: input ids (B, max_seq)
        :type data: Numpy array
        :param token_id2idx: mapping to convert bert token to idx (which goes from 0 to V-1)
        '''
        self.data = data[:,:-1].copy() # (B, max_seq -1)
        self.labels = data[:, 1:].copy() # (B, max_seq -1)
        assert self.data.shape == self.labels.shape, "data shape and label shape is different, this will cause error during TBPTT"
        for r in range(self.labels.shape[0]):
            for t in range(self.labels.shape[1]):
              self.labels[r, t] = token_id2idx[self.labels[r, t]]
        self.max_seq = self.data.shape[1]

    def __len__(self):
        'Denotes the total number of samples'
        return self.data.shape[0]

    def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        X = self.data[index] # (max_seq)
        y = self.labels[index] # (max_seq - 1)
        # y = torch.ones(1).fill_(self.label)
        return X, y

Environment
PyTorch version: 1.3.1+cu92
Is debug build: No
CUDA used to build PyTorch: 9.2.148
OS: Debian GNU/Linux 9.6 (stretch)
GCC version: (Debian 6.3.0-18+deb9u1) 6.3.0 20170516
CMake version: Could not collect
Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: Tesla K80
Nvidia driver version: 410.72
cuDNN version: Could not collect
Versions of relevant libraries:
[pip] numpy==1.16.4
[pip] numpydoc==0.9.1
[pip] pytorch-ignite==0.1.0
[pip] pytorch-lightning==0.6.0
[pip] pytorch-transformers==1.1.0
[pip] torch==1.3.1+cu92
[pip] torchfile==0.1.0
[pip] torchvision==0.4.2+cu92
[conda] blas                      1.0                         mkl
[conda] mkl                       2019.4                      243
[conda] mkl-service               2.0.2            py37h7b6447c_0
[conda] mkl_fft                   1.0.12           py37ha843d7b_0
[conda] mkl_random                1.0.2            py37hd81dba3_0
[conda] pytorch-ignite            0.1.0                    pypi_0    pypi
[conda] pytorch-lightning         0.6.0                    pypi_0    pypi
[conda] pytorch-transformers      1.1.0                    pypi_0    pypi
[conda] torch                     1.3.1+cu92               pypi_0    pypi
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchvision               0.4.2+cu92               pypi_0    pypi