Megatron-LM checkpoints don't work properly
I don't think this command
```
python -m pip install --upgrade git+https://github.com/NVIDIA/NeMo.git@main
``` 
is installing NeMo properly. The error you are getting is from using an old version of NeMo with a new version of Megatron-LM. When we upgrade Megatron-LM, we fixed that error. This has all happened post 1.0.0b1 release.

Try:
```
git clone https://github.com/NVIDIA/NeMo.git
cd NeMo
sh reinstall.sh
``` Please let us know if the above suggestion was helpful, thanks
**Describe the bug**

Megatron-LM checkpoints appear to not work properly with the recent changes in NLP module.

**Steps/Code to reproduce bug**

1. implement minimal change to the [Token Classification yaml file](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/conf/token_classification_config.yaml) to use Meagtron-LM checkpoint.
  1.1. Change [`bert-base-uncased`](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/conf/token_classification_config.yaml#L77) to `megatron-bert-345m-cased`
  1.2. Download and unpack the `gmb_v-2.2.0_clean` dataset for NER as instructed in the [tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Token_Classification_Named_Entity_Recognition.ipynb), specify the [location of the dataset in the yaml file](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/token_classification/conf/token_classification_config.yaml#L44).
2. Run token classification example `python token_classification.py`


**Expected behavior**

Instead of training, error occurs and the program exits:
`TypeError: bert_extended_attention_mask() missing 1 required positional argument: 'dtype`

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

 - Environment location: Based on [NeMo NGC container](nvcr.io/nvidia/nemo:v1.0.0b1), upgrade NeMo with pip install: `python -m pip install --upgrade git+https://github.com/NVIDIA/NeMo.git@main`
 Because without the additional pip install, different error occurs (another bug - mismatch between NGC container and pip package).
 - Method of NeMo install: pip install (`python -m pip install --upgrade git+https://github.com/NVIDIA/NeMo.git@main`)
 - If method of install is [Docker], provide `docker pull` & `docker run` commands used: `docker pull nvcr.io/nvidia/nemo:v1.0.0b1`

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version: `nvcr.io/nvidia/nemo:v1.0.0b1` docker container
- PyTorch version: `nvcr.io/nvidia/nemo:v1.0.0b1` docker container
- Python version: `nvcr.io/nvidia/nemo:v1.0.0b1` docker container

