Conformer model fails to load on CPU after export
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

Conformer model cannot be loaded when loaded on CPU.
There an unclear `RuntimeError`: 
```
Traceback (most recent call last):
  File "bug_nemo.py", line 12, in <module>
    ts_model = torch.jit.load(f"/tmp/{model_name}.pt")
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_serialization.py", line 162, in load
    cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files)
RuntimeError: required keyword attribute 'value' is undefined
```

It works fine when loading on GPU or using Citrinet, for example : `stt_en_citrinet_1024`.

**Steps/Code to reproduce bug**

```python:bug_nemo.py
import torch
import nemo.collections.asr as nemo_asr

model_name = "stt_en_conformer_ctc_small"

nemo_model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(
    model_name=model_name,
    map_location="cpu"
)

nemo_model.export(f"/tmp/{model_name}.pt")

ts_model = torch.jit.load(f"/tmp/{model_name}.pt")
```
Output:
```
/opt/conda/lib/python3.8/site-packages/apex/pyprof/__init__.py:5: FutureWarning: pyprof will be removed by the end of June, 2022
  warnings.warn("pyprof will be removed by the end of June, 2022", FutureWarning)
[NeMo I 2022-08-30 22:17:01 cloud:56] Found existing object /home/tarteel/.cache/torch/NeMo/NeMo_1.10.0/stt_en_conformer_ctc_small/5d2d8e5b2b5adb8f5091363c6ba19c55/stt_en_conformer_ctc_small.nemo.
[NeMo I 2022-08-30 22:17:01 cloud:62] Re-using file from: /home/tarteel/.cache/torch/NeMo/NeMo_1.10.0/stt_en_conformer_ctc_small/5d2d8e5b2b5adb8f5091363c6ba19c55/stt_en_conformer_ctc_small.nemo
[NeMo I 2022-08-30 22:17:01 common:789] Instantiating model from pre-trained checkpoint
[NeMo I 2022-08-30 22:17:02 mixins:166] Tokenizer SentencePieceTokenizer initialized with 1024 tokens
[NeMo W 2022-08-30 22:17:02 modelPT:148] If you intend to do training or fine-tuning, please call the ModelPT.setup_training_data() method and provide a valid configuration file to setup the train data loader.
    Train config : 
    manifest_filepath: /data/NeMo_ASR_SET/English/v2.0/train/tarred_audio_manifest.json
    sample_rate: 16000
    batch_size: 64
    shuffle: true
    num_workers: 8
    pin_memory: true
    use_start_end_token: false
    trim_silence: false
    max_duration: 20.0
    min_duration: 0.1
    shuffle_n: 2048
    is_tarred: true
    tarred_audio_filepaths: /data/NeMo_ASR_SET/English/v2.0/train/audio__OP_0..4095_CL_.tar
    
[NeMo W 2022-08-30 22:17:02 modelPT:155] If you intend to do validation, please call the ModelPT.setup_validation_data() or ModelPT.setup_multiple_validation_data() method and provide a valid configuration file to setup the validation data loader(s). 
    Validation config : 
    manifest_filepath:
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-dev-other.json
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-dev-clean.json
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-test-other.json
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-test-clean.json
    sample_rate: 16000
    batch_size: 64
    shuffle: false
    num_workers: 8
    pin_memory: true
    use_start_end_token: false
    is_tarred: false
    tarred_audio_filepaths: na
    
[NeMo W 2022-08-30 22:17:02 modelPT:161] Please call the ModelPT.setup_test_data() or ModelPT.setup_multiple_test_data() method and provide a valid configuration file to setup the test data loader(s).
    Test config : 
    manifest_filepath:
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-test-other.json
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-dev-clean.json
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-dev-other.json
    - /data/ASR/LibriSpeech/librispeech_withsp2/manifests/librivox-test-clean.json
    sample_rate: 16000
    batch_size: 64
    shuffle: false
    num_workers: 8
    pin_memory: true
    use_start_end_token: false
    is_tarred: false
    tarred_audio_filepaths: na
    
[NeMo I 2022-08-30 22:17:02 features:200] PADDING: 0
[NeMo W 2022-08-30 22:17:03 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torchmetrics/utilities/prints.py:36: UserWarning: Torchmetrics v0.9 introduced a new argument class property called `full_state_update` that has
                    not been set for this class (WER). The property determines if `update` by
                    default needs access to the full metric state. If this is not the case, significant speedups can be
                    achieved and we recommend setting this to `False`.
                    We provide an checking function
                    `from torchmetrics.utilities import check_forward_no_full_state`
                    that can be used to check if the `full_state_update=True` (old and potential slower behaviour,
                    default for now) or if `full_state_update=False` can be used safely.
                    
      warnings.warn(*args, **kwargs)
    
[NeMo W 2022-08-30 22:17:03 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torchmetrics/utilities/prints.py:36: UserWarning: Torchmetrics v0.9 introduced a new argument class property called `full_state_update` that has
                    not been set for this class (WERBPE). The property determines if `update` by
                    default needs access to the full metric state. If this is not the case, significant speedups can be
                    achieved and we recommend setting this to `False`.
                    We provide an checking function
                    `from torchmetrics.utilities import check_forward_no_full_state`
                    that can be used to check if the `full_state_update=True` (old and potential slower behaviour,
                    default for now) or if `full_state_update=False` can be used safely.
                    
      warnings.warn(*args, **kwargs)
    
[NeMo I 2022-08-30 22:17:03 save_restore_connector:243] Model EncDecCTCModelBPE was successfully restored from /home/tarteel/.cache/torch/NeMo/NeMo_1.10.0/stt_en_conformer_ctc_small/5d2d8e5b2b5adb8f5091363c6ba19c55/stt_en_conformer_ctc_small.nemo.
[NeMo I 2022-08-30 22:17:03 export_utils:263] Swapped 96 modules
[NeMo W 2022-08-30 22:17:04 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/_jit_internal.py:669: LightningDeprecationWarning: The `LightningModule.model_size` property was deprecated in v1.5 and will be removed in v1.7. Please use the `pytorch_lightning.utilities.memory.get_model_size_mb`.
      if hasattr(mod, name):
    
[NeMo W 2022-08-30 22:17:04 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1402: UserWarning: positional arguments and argument "destination" are deprecated. nn.Module.state_dict will not accept them in the future. Refer to https://pytorch.org/docs/master/generated/torch.nn.Module.html#torch.nn.Module.state_dict for details.
      warnings.warn(
    
[NeMo W 2022-08-30 22:17:04 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py:1402: UserWarning: argument "destination" are deprecated. nn.Module.state_dict will not accept them in the future. Refer to https://pytorch.org/docs/master/generated/torch.nn.Module.html#torch.nn.Module.state_dict for details.
      warnings.warn(
    
[NeMo W 2022-08-30 22:17:04 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/_jit_internal.py:670: LightningDeprecationWarning: The `LightningModule.model_size` property was deprecated in v1.5 and will be removed in v1.7. Please use the `pytorch_lightning.utilities.memory.get_model_size_mb`.
      item = getattr(mod, name)
    
[NeMo W 2022-08-30 22:17:04 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/_jit_internal.py:669: LightningDeprecationWarning: `LightningModule.use_amp` was deprecated in v1.6 and will be removed in v1.8. Please use `Trainer.amp_backend`.
      if hasattr(mod, name):
    
[NeMo W 2022-08-30 22:17:04 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/torch/_jit_internal.py:670: LightningDeprecationWarning: `LightningModule.use_amp` was deprecated in v1.6 and will be removed in v1.8. Please use `Trainer.amp_backend`.
      item = getattr(mod, name)
    
[NeMo W 2022-08-30 22:17:04 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/nemo/collections/asr/modules/conformer_encoder.py:248: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
      if max_audio_length > self.max_audio_length:
    
[NeMo I 2022-08-30 22:17:06 exportable:84] Successfully exported EncDecCTCModelBPE to /tmp/stt_en_conformer_ctc_small.pt
Traceback (most recent call last):
  File "bug_nemo.py", line 12, in <module>
    ts_model = torch.jit.load(f"/tmp/{model_name}.pt")
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_serialization.py", line 162, in load
    cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files)
RuntimeError: required keyword attribute 'value' is undefined
```

**Expected behavior**

Export a CPU TorchScript model.

If `map_location=cpu` is removed, the exported GPU model works. 

I’ve tried moving to CPU after loading it, with no success:
```
nemo_model = nemo_model.to(‘cpu’)
```



**Environment details**

Docker image: `nvcr.io/nvidia/nemo:22.05`