RuntimeError: expected scalar type Half but found Float
State: closed | Pull request: no

I have trained nemo 1.8.2 model, earlier I have used nemo 1.4.0 and used below command to classify the texts

model.classifytext(queries=config.model.infer_samples, batch_size=8, max_seq_length=511)

This time when calling the above command with nemo 1.8.2 I am facing issue as **RuntimeError: expected scalar type Half but found Float**. 

Below is the entire log trace:
RuntimeError                              Traceback (most recent call last)
Input In [20], in <cell line: 1>()
----> 1 model.classifytext(queries=config.model.infer_samples, batch_size=8, max_seq_length=511)

File /opt/conda/lib/python3.8/site-packages/torch/autograd/grad_mode.py:27, in _DecoratorContextManager.__call__.<locals>.decorate_context(*args, **kwargs)
     24 @functools.wraps(func)
     25 def decorate_context(*args, **kwargs):
     26     with self.clone():
---> 27         return func(*args, **kwargs)

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/text_classification/text_classification_model.py:259, in TextClassificationModel.classifytext(self, queries, batch_size, max_seq_length)
    256 for i, batch in enumerate(infer_datalayer):
    257     input_ids, input_type_ids, input_mask, subtokens_mask = batch
--> 259     logits = self.forward(
    260         input_ids=input_ids.to(device),
    261         token_type_ids=input_type_ids.to(device),
    262         attention_mask=input_mask.to(device),
    263     )
    265     preds = tensor2list(torch.argmax(logits, axis=-1))
    266     all_preds.extend(preds)

File /opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py:888, in typecheck.__call__(self, wrapped, instance, args, kwargs)
    886 # If types are not defined, skip type checks and just call the wrapped method
    887 if input_types is None and output_types is None:
--> 888     return wrapped(*args, **kwargs)
    890 # Check that all arguments are kwargs
    891 if input_types is not None and len(args) > 0:

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/text_classification/text_classification_model.py:82, in TextClassificationModel.forward(self, input_ids, attention_mask, token_type_ids)
     76 @typecheck()
     77 def forward(self, input_ids, attention_mask, token_type_ids):
     78     """
     79     No special modification required for Lightning, define it as you normally would
     80     in the `nn.Module` in vanilla PyTorch.
     81     """
---> 82     hidden_states = self.bert_model(
     83         input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask
     84     )
     85     if isinstance(hidden_states, tuple):
     86         hidden_states = hidden_states[0]

File /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1111, in Module._call_impl(self, *input, **kwargs)
   1107 # If we don't have any hooks, we want to skip the rest of the logic in
   1108 # this function, and just call forward.
   1109 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1110         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1111     return forward_call(*input, **kwargs)
   1112 # Do not call functions when jit is used
   1113 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron_bert_model.py:130, in MegatronBertModel.forward(self, input_ids, attention_mask, token_type_ids, lm_labels)
    129 def forward(self, input_ids, attention_mask, token_type_ids, lm_labels=None):
--> 130     output_tensor = self.model(input_ids, attention_mask, token_type_ids=token_type_ids, lm_labels=lm_labels)
    131     return output_tensor

File /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1111, in Module._call_impl(self, *input, **kwargs)
   1107 # If we don't have any hooks, we want to skip the rest of the logic in
   1108 # this function, and just call forward.
   1109 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1110         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1111     return forward_call(*input, **kwargs)
   1112 # Do not call functions when jit is used
   1113 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/models/language_modeling/megatron/bert_model.py:226, in BertModel.forward(self, bert_model_input, attention_mask, token_type_ids, lm_labels)
    223 input_ids = bert_model_input
    224 position_ids = build_position_ids(input_ids)
--> 226 lm_output = self.language_model(
    227     input_ids, position_ids, extended_attention_mask, token_type_ids=token_type_ids
    228 )
    230 if self.post_process and self.add_binary_head:
    231     lm_output, pooled_output = lm_output

File /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1111, in Module._call_impl(self, *input, **kwargs)
   1107 # If we don't have any hooks, we want to skip the rest of the logic in
   1108 # this function, and just call forward.
   1109 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1110         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1111     return forward_call(*input, **kwargs)
   1112 # Do not call functions when jit is used
   1113 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/language_model.py:746, in TransformerLanguageModel.forward(self, enc_input_ids, enc_position_ids, enc_attn_mask, prompt_ids, dec_input_ids, dec_position_ids, dec_attn_mask, enc_dec_attn_mask, token_type_ids, layer_past, get_key_value, pooling_sequence_index, enc_hidden_states, output_enc_hidden_only, encoder_input, set_inference_key_value_memory, inference_max_sequence_len)
    744 # encoder.
    745 if enc_hidden_states is None:
--> 746     encoder_output = self.encoder(
    747         encoder_input,
    748         enc_attn_mask,
    749         layer_past=layer_past,
    750         get_key_value=get_key_value,
    751         set_inference_key_value_memory=set_inference_key_value_memory,
    752         inference_max_sequence_len=inference_max_sequence_len,
    753     )
    754 else:
    755     encoder_output = enc_hidden_states.to(encoder_input.dtype)

File /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1111, in Module._call_impl(self, *input, **kwargs)
   1107 # If we don't have any hooks, we want to skip the rest of the logic in
   1108 # this function, and just call forward.
   1109 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1110         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1111     return forward_call(*input, **kwargs)
   1112 # Do not call functions when jit is used
   1113 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/lib/python3.8/site-packages/nemo/collections/nlp/modules/common/megatron/transformer.py:1095, in ParallelTransformer.forward(self, hidden_states, attention_mask, layer_past, get_key_value, encoder_output, enc_dec_attn_mask, set_inference_key_value_memory, inference_max_sequence_len)
   1092 if self.post_process:
   1093     # Reverting data format change [s b h] --> [b s h].
   1094     hidden_states = hidden_states.transpose(0, 1).contiguous()
-> 1095     output = self.final_layernorm(hidden_states)
   1096 else:
   1097     output = hidden_states

File /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1111, in Module._call_impl(self, *input, **kwargs)
   1107 # If we don't have any hooks, we want to skip the rest of the logic in
   1108 # this function, and just call forward.
   1109 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1110         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1111     return forward_call(*input, **kwargs)
   1112 # Do not call functions when jit is used
   1113 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/lib/python3.8/site-packages/apex/normalization/fused_layer_norm.py:414, in MixedFusedLayerNorm.forward(self, input)
    412 if not input.is_cuda:
    413     return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps)
--> 414 return mixed_dtype_fused_layer_norm_affine(input, self.weight, self.bias, self.normalized_shape, self.eps)

File /opt/conda/lib/python3.8/site-packages/apex/normalization/fused_layer_norm.py:183, in mixed_dtype_fused_layer_norm_affine(input, weight, bias, normalized_shape, eps)
    181 args = _cast_if_autocast_enabled(input, weight, bias, normalized_shape, eps)
    182 with torch.cuda.amp.autocast(enabled=False):
--> 183     return FusedLayerNormAffineMixedDtypesFunction.apply(*args)

File /opt/conda/lib/python3.8/site-packages/apex/normalization/fused_layer_norm.py:96, in FusedLayerNormAffineMixedDtypesFunction.forward(ctx, input, weight, bias, normalized_shape, eps)
     94 weight_ = weight.contiguous()
     95 bias_ = bias.contiguous()
---> 96 output, mean, invvar = fused_layer_norm_cuda.forward_affine_mixed_dtypes(
     97     input_, ctx.normalized_shape, weight_, bias_, ctx.eps
     98 )
     99 ctx.save_for_backward(input_, weight_, bias_, mean, invvar)
    100 return output

RuntimeError: expected scalar type Half but found Float