Unable to train Entity Linking
I am trying to train Entity Linking in NLP. 

It runs fine. However, while saving results after Epoch 0 it gives this error 
In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present

I checked and training step method does have a dictionary. What could be the reason for this?
Im running training on Google Colab (if that helps )