RuntimeError raised when using classifytext function from text_classification_class.py with a pretrained model using tensor model parallel
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**
We have a BERT model pretrained using Megatron-LM (v2.2) with the tensor model parallel set to 2.
We use NeMo to fine-tune the model for text classification.
We do not have an issue with fine-tuning and test, but we have an issue with prediction.
We try to apply the classifytext function from the text_classification_class.py to get the predicted labels, but it raised a RuntimeError as Tensors must be CUDA and dense (detailed error trace is attached below).

**Steps/Code to reproduce bug**
Following is the code snippet we used (similar as the tutorial)
```python
trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **config.trainer)
exp_dir = exp_manager(trainer, config.get("exp_manager", None))
model = nemo_nlp.models.TextClassificationModel(cfg=config.model, trainer=trainer)
trainer.fit(model)
quries = load_test(config.model.test_ds.file_path)
all_preds = model.classifytext(quries, config.model.test_ds.batch_size, config.model.dataset.max_seq_length)
```

**Expected behavior**
Yield predicted labels.

**actual behavior**
raised a RuntimeError as
```
Traceback (most recent call last):
  File "/red/gatortron-phi/workspace/2021gatortron/nemo_downstream/re/re_ddp.py", line 66, in <module>
    main(global_args)
  File "/red/gatortron-phi/workspace/2021gatortron/nemo_downstream/re/re_ddp.py", line 46, in main
    all_preds = model.classifytext(config.model.infer_samples, config.model.test_ds.batch_size, config.model.dataset.max_seq_length)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/text_classification/text_classification_model.py", line 280, in classifytext
    logits = self.forward(
  File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py", line 792, in __call__
    outputs = wrapped(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/text_classification/text_classification_model.py", line 105, in forward
    hidden_states = self.bert_model(
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1015, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py", line 792, in __call__
    outputs = wrapped(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/megatron_bert.py", line 182, in forward
    sequence_output = self.language_model(
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1015, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py", line 414, in forward
    return super(TransformerLanguageModel, self).forward(
  File "/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py", line 321, in forward
    embedding_output = self.embedding(input_ids, position_ids,
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1015, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/megatron/model/language_model.py", line 183, in forward
    words_embeddings = self.word_embeddings(input_ids)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1015, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/megatron/mpu/layers.py", line 191, in forward
    output = reduce_from_tensor_model_parallel_region(output_parallel)
  File "/opt/conda/lib/python3.8/site-packages/megatron/mpu/mappings.py", line 149, in reduce_from_tensor_model_parallel_region
    return _ReduceFromModelParallelRegion.apply(input_)
  File "/opt/conda/lib/python3.8/site-packages/megatron/mpu/mappings.py", line 101, in forward
    return _reduce(input_)
  File "/opt/conda/lib/python3.8/site-packages/megatron/mpu/mappings.py", line 30, in _reduce
    torch.distributed.all_reduce(input_, group=get_tensor_model_parallel_group())
  File "/opt/conda/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 1175, in all_reduce
    work = group.allreduce([tensor], opts)
RuntimeError: Tensors must be CUDA and dense
```

**Environment overview (please complete the following information)**
 - Environment location: Slurm, 1 node 8 GPUs, NeMo-1.2.0 docker container from NGC