Hybrid decoder export

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>
borisfom · Jul 5, 2023 · cd6d11d · cd6d11d
1 parent 8d78af3
commit cd6d11d
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/nemo/collections/asr/models/hybrid_rnnt_ctc_models.py b/nemo/collections/asr/models/hybrid_rnnt_ctc_models.py
@@ -650,7 +650,14 @@ def list_export_subnets(self):
  if self.cur_decoder == 'rnnt':
  return ['encoder', 'decoder_joint']
  else:
- return ['encoder']
+ return ['self']
+
+ @property
+ def output_module(self):
+ if self.cur_decoder == 'rnnt':
+ return self.decoder
+ else:
+ return self.ctc_decoder
 
  @classmethod
  def list_available_models(cls) -> Optional[PretrainedModelInfo]: