fixed the onnx bug in conformer for non-streaming models. (#5242)

Signed-off-by: Vahid <vnoroozi@nvidia.com> Signed-off-by: Vahid <vnoroozi@nvidia.com> Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>
NVIDIA · Nov 17, 2022 · 1d9488a · 1d9488a
1 parent 988dedb
commit 1d9488a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/collections/asr/modules/conformer_encoder.py b/nemo/collections/asr/modules/conformer_encoder.py
@@ -99,10 +99,10 @@ def input_example(self, max_batch=1, max_dim=256):
         if hasattr(self, 'export_cache_support') and self.export_cache_support:
             cache_last_channel = torch.randn(self.n_layers, max_batch, max_dim, self.d_model).to(dev)
             cache_last_time = torch.randn(self.n_layers, max_batch, self.d_model, self.conv_context_size[0]).to(dev)
+            all_input_example = tuple([input_example, input_example_length, cache_last_channel, cache_last_time])
         else:
-            cache_last_channel = cache_last_time = None
+            all_input_example = tuple([input_example, input_example_length])
 
-        all_input_example = tuple([input_example, input_example_length, cache_last_channel, cache_last_time])
         return all_input_example
 
     @property