Error while trying to input tokenized text into NeMo's MegatronBertModel
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

Hello, I am trying to load a `.nemo` checkpoint file into NeMo's `MegatronBertModel()` instance using the `restore_from()` method as shown below. The tokenizer file used is the BertWordPiece tokenizer pretrained for this task. My code is as shown below. 

```python
from nemo.collections.nlp.models.language_modeling.megatron_bert_model import MegatronBertModel
from transformers import BertTokenizer

model = MegatronBertModel.restore_from("MegatronBERT.nemo", trainer=trainer).to(device)

tokenizer=BertTokenizer("8b2c88a36abf4b40ba52976353202740_vocab.txt")

sentence = "I am not feeling well"

input = tokenizer(sentence)

input["input_ids"] = torch.LongTensor(input["input_ids"]).to(device)

input["attention_mask"] = torch.FloatTensor(input["attention_mask"]).to(device)

input["token_type_ids"] = torch.LongTensor(input["token_type_ids"]).to(device)

model.eval()

model(**input)

``` 

I convert the `input_ids`, `attention_mask` and the `token_type_ids` into tensors since they are returned as lists from the output of the tokenizer & if I don't do so, it gives an error asking me to. On executing the above code, I run into the below error

```python
Error executing job with overrides: []
Traceback (most recent call last):
  File "megatron_bert_eval.py", line 159, in <module>
    main()
  File "/opt/conda/lib/python3.8/site-packages/nemo/core/config/hydra_runner.py", line 104, in wrapper
    _run_hydra(
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 377, in _run_hydra
    run_and_report(
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 214, in run_and_report
    raise ex
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 211, in run_and_report
    return func()
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 378, in <lambda>
    lambda: hydra.run(
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/hydra.py", line 111, in run
    _ = ret.return_value
  File "/opt/conda/lib/python3.8/site-packages/hydra/core/utils.py", line 233, in return_value
    raise self._return_value
  File "/opt/conda/lib/python3.8/site-packages/hydra/core/utils.py", line 160, in run_job
    ret.return_value = task_function(task_cfg)
  File "megatron_bert_eval.py", line 93, in main
    model(**input)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron_bert_model.py", line 98, in forward
    output_tensor = self.model(input_ids, attention_mask, token_type_ids=token_type_ids, lm_labels=lm_labels)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron/bert_model.py", line 222, in forward
    extended_attention_mask = bert_extended_attention_mask(attention_mask)
  File "/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron/bert_model.py", line 51, in bert_extended_attention_mask
    attention_mask_bs1 = attention_mask.unsqueeze(2)
IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)
```

Can somebody kindly help me with this error?

Thanks in advance!

**Environment overview (please complete the following information)**

 - Environment: I am using `nvcr.io/nvidia/nemo:22.09` docker image
 - Method of NeMo install: I cloned the NeMo branch `1.10.0`  in my docker image