Issue in loading model from checkpoint
**Describe the bug**

I am trying to run the Relation Extraction - BioMegatron notebook where once the model is trained, I am trying to load the best checkpoint.

**Steps/Code to reproduce bug**

- Train the model as usual

`checkpoint_path = trainer.checkpoint_callback.best_model_path`
`eval_model = nemo_nlp.models.TextClassificationModel.load_from_checkpoint(checkpoint_path=checkpoint_path)`



**Expected behavior**

<img width="1048" alt="Screen Shot 2021-01-25 at 8 01 09 PM" src="https://user-images.githubusercontent.com/7240751/105785078-13d5e100-5f48-11eb-991d-3560fa24a1e3.png">


**Environment overview (please complete the following information)**

 - Environment Location: Collab
 - Method of NeMo install:  pip 


