Fix torch.jit.script() export for pruned_transducer_stateless2 (#1410)

k2-fsa · Dec 10, 2023 · b0f70c9 · b0f70c9
1 parent df56aff
commit b0f70c9
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 0 deletions.
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/export.py b/egs/librispeech/ASR/pruned_transducer_stateless2/export.py
@@ -49,6 +49,7 @@
 
 import k2
 import torch
+from scaling_converter import convert_scaled_to_non_scaled
 from train import add_model_arguments, get_params, get_transducer_model
 
 from icefall.checkpoint import average_checkpoints, find_checkpoints, load_checkpoint
@@ -198,6 +199,7 @@ def main():
     model.eval()
 
     if params.jit:
+        convert_scaled_to_non_scaled(model, inplace=True)
         # We won't use the forward() method of the model in C++, so just ignore
         # it here.
         # Otherwise, one of its arguments is a ragged tensor and is not

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/lstmp.py b/egs/librispeech/ASR/pruned_transducer_stateless2/lstmp.py
@@ -0,0 +1 @@
+../lstm_transducer_stateless2/lstmp.py
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless2/scaling_converter.py b/egs/librispeech/ASR/pruned_transducer_stateless2/scaling_converter.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless3/scaling_converter.py
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		../pruned_transducer_stateless3/scaling_converter.py