export Titanet-l.nemo to onnx get shape error
State: closed | Pull request: no
Labels: name: bug

Hi there,
     I use [export.py](https://github.com/NVIDIA/NeMo/blob/6dd4263fecd03a39688dfd99866ac8d5842b65e6/scripts/export.py) to export titanet to onnx, and succeed(get titanet-l.onnx). Then I use tensorRT OnnxParser to parse this onnx, and get error like:
```
[03/22/2022-09:41:57] [TRT] [E] ModelImporter.cpp:774: --- Begin node ---
[03/22/2022-09:41:57] [TRT] [E] ModelImporter.cpp:775: input: "608"
output: "609"
name: "If_380"
op_type: "If"
attribute {
  name: "then_branch"
  g {
    node {
      output: "610"
      name: "Constant_381"
      op_type: "Constant"
      attribute {
        name: "value"
        t {
          dims: 1
          data_type: 7
          raw_data: "\377\377\377\377\377\377\377\377"
        }
        type: TENSOR
      }
      doc_string: ""
    }
    node {
      input: "601"
      input: "610"
      output: "611"
      name: "Squeeze_382"
      op_type: "Squeeze"
      doc_string: ""
    }
    name: "torch-jit-export1"
    output {
      name: "611"
      type {
        tensor_type {
          elem_type: 1
          shape {
            dim {
              dim_param: "Squeeze611_dim_0"
            }
            dim {
              dim_param: "Squeeze611_dim_1"
            }
          }
        }
      }
    }
  }
  type: GRAPH
}
attribute {
  name: "else_branch"
  g {
    node {
      input: "601"
      output: "612"
      name: "Identity_383"
      op_type: "Identity"
      doc_string: ""
    }
    name: "torch-jit-export2"
    output {
      name: "612"
      type {
        tensor_type {
          elem_type: 1
          shape {
            dim {
              dim_param: "Identity612_dim_0"
            }
            dim {
              dim_param: "Identity612_dim_1"
            }
            dim {
              dim_param: "Identity612_dim_2"
            }
          }
        }
      }
    }
  }
  type: GRAPH
}
doc_string: "/root/miniconda3/envs/myconda/lib/python3.8/site-packages/nemo/collections/asr/modules/conv_asr.py(844): forward\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/nn/modules/module.py(1090): _slow_forward\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/nemo/collections/asr/models/label_models.py(235): forward_for_export\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/nn/modules/module.py(1090): _slow_forward\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/jit/_trace.py(118): wrapper\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/jit/_trace.py(127): forward\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/nn/modules/module.py(1102): _call_impl\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/jit/_trace.py(1166): _get_trace_graph\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/onnx/utils.py(388): _trace_and_get_graph_from_model\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/onnx/utils.py(437): _create_jit_graph\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/onnx/utils.py(493): _model_to_graph\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/onnx/utils.py(724): _export\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/onnx/utils.py(107): export\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/torch/onnx/__init__.py(316): export\n/root/miniconda3/envs/myconda/lib/python3.8/site-packages/nemo/core/classes/exportable.py(142): export\n/home/ccl/sv/test/export.py(137): nemo_export\n/home/ccl/sv/test/export.py(162): <module>\n/root/ccl/.pycharm_helpers/pydev/_pydev_imps/_pydev_execfile.py(18): execfile\n/root/ccl/.pycharm_helpers/pydev/pydevd.py(1147): run\n/root/ccl/.pycharm_helpers/pydev/pydevd.py(1752): main\n/root/ccl/.pycharm_helpers/pydev/pydevd.py(1758): <module>\n"

[03/22/2022-09:41:57] [TRT] [E] ModelImporter.cpp:776: --- End node ---
[03/22/2022-09:41:57] [TRT] [E] ModelImporter.cpp:779: ERROR: ModelImporter.cpp:179 In function parseGraph:
[6] Invalid Node - If_380
If_380_OutputLayer: IIfConditionalOutputLayer inputs must have the same shape.
```

Here is my parse code:

```     
    emb_onnx_path = 'titanet-l.onnx'
    logger = trt.Logger(trt.Logger.WARNING)
    network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    with trt.Builder(logger) as builder:
        with builder.create_network(flags=network_flags) as network, trt.OnnxParser(
                network, logger
        ) as parser, builder.create_builder_config() as builder_config:
            parser.parse_from_file(emb_onnx_path)
```

I don't know whether it's a bug or just usage error. Could you please have a look?