Fix torchscript export error of FastPitchModel
State: closed | Pull request: yes
Labels: name: TTS

name: stale

# What does this PR do?

When trying to export FastPitchModel to TorchScript model:
```python
from nemo.collections.tts.models import FastPitchModel

fastpitch_model = FastPitchModel.from_pretrained("tts_en_fastpitch")
fastpitch_model.export("fastpitch.pt")
```
It reported error:
```
Traceback (most recent call last):
  File "/home/robin/try1/NeMo/export.py", line 11, in <module>
    fastpitch_model.export("fastpitch.pt")
  File "/home/robin/try1/NeMo/nemo/core/classes/exportable.py", line 113, in export
    out, descr, out_example = model._export(
  File "/home/robin/try1/NeMo/nemo/core/classes/exportable.py", line 200, in _export
    jitted_model = torch.jit.trace_module(
  File "/home/robin/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/jit/_trace.py", line 1065, in trace_module
    module._c._create_method_from_trace(
  File "/home/robin/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/robin/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/robin/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1508, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/robin/try1/NeMo/nemo/collections/tts/models/fastpitch.py", line 861, in forward_for_export
    return self.fastpitch.infer(text=text, pitch=pitch, pace=pace, volume=volume, speaker=speaker)
  File "/home/robin/try1/NeMo/nemo/collections/tts/modules/fastpitch.py", line 430, in infer
    volume_extended, _ = regulate_len(durs_predicted, volume.unsqueeze(-1), pace)
  File "/home/robin/try1/NeMo/nemo/collections/tts/parts/utils/helpers.py", line 597, in regulate_len
    enc_rep = torch.matmul(mult, enc_out)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (270x44 and 1x1)
```
The `torch.jit.trace_module()` in `nemo/core/classes/exportable.py` could only process tuple of tensors. The `input_examle` of FastPitchModel provides four arguments `text`/`pitch`/`pace`/`batch_lengths` but the first four arguments `forward_for_export()` of FastPitchModel are `text`/`pitch`/`pace`/`volume` (the fifth one is `batch_lengths`). Then the `torch.jit.trace_module()` will mistakenly recognize `batch_lengths` input as `volume` argument.
A simple solution for this is swap arguments of `volume` and `batch_lengths` in `forward_for_export()`.

After fix the input, another error jump out:
```
Traceback (most recent call last):
  File "/home/robin/try1/NeMo/export.py", line 5, in <module>
    fastpitch_model.export("fastpitch.pt")
  File "/home/robin/try1/NeMo/nemo/core/classes/exportable.py", line 113, in export
    out, descr, out_example = model._export(
  File "/home/robin/try1/NeMo/nemo/core/classes/exportable.py", line 203, in _export
    jitted_model = torch.jit.trace_module(
  File "/home/robin/miniconda3/envs/nemo/lib/python3.10/site-packages/torch/jit/_trace.py", line 1065, in trace_module
    module._c._create_method_from_trace(
RuntimeError: Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions
```
The reason is that the output of `infer()` for `FastPitchModule` contains a `None` argument `volume_extended`. To avoid this, we shall check before return output.

# Changelog 
- Fix torchscript export error of FastPitchModel

# Usage
After this PR, we can export both ONNX and TorchScript successfully

```python
from nemo.collections.tts.models import FastPitchModel

fastpitch_model = FastPitchModel.from_pretrained("tts_en_fastpitch")
fastpitch_model.export("fastpitch.onnx")
fastpitch_model.export("fastpitch.pt")
```

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [x] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

@XuesongYang @blisc @okuchaiev