Specify label name
I'm trying to train a classification network on T5 encoder from this training code:
https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb
As I have a special output for label, I need to make trainer associate field from batch with field from my net output dictionary, as they represent output class label. How to do that?