BERT model throws error when used in Pytorch Lightning

[goku]

I think some mismatch with positional arguments is there:

try:

@sudarshan85

def forward(self, input_ids, attention_mask, token_type_ids, labels):
    return self.model(input_ids, attention_mask, token_type_ids, labels=labels)

or maybe

@sudarshan85

 def forward(self, *args, **kwargs):
    return self.model(*args, **kwargs)