From 0af3e7beda1cb47cba8b51ce71f691e86cae2091 Mon Sep 17 00:00:00 2001
From: Fangjun Kuang <csukuangfj@gmail.com>
Date: Mon, 16 Jan 2023 20:26:36 +0800
Subject: [PATCH] fix export for stateless4 (#844)

---
 egs/librispeech/ASR/pruned_transducer_stateless4/export.py      | 2 ++
 egs/librispeech/ASR/pruned_transducer_stateless4/lstmp.py       | 1 +
 .../ASR/pruned_transducer_stateless4/scaling_converter.py       | 1 +
 3 files changed, 4 insertions(+)
 create mode 120000 egs/librispeech/ASR/pruned_transducer_stateless4/lstmp.py
 create mode 120000 egs/librispeech/ASR/pruned_transducer_stateless4/scaling_converter.py

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless4/export.py b/egs/librispeech/ASR/pruned_transducer_stateless4/export.py
index 401b3ef3a..8f33f5b05 100755
--- a/egs/librispeech/ASR/pruned_transducer_stateless4/export.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless4/export.py
@@ -50,6 +50,7 @@ from pathlib import Path
 
 import sentencepiece as spm
 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 (
@@ -261,6 +262,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_stateless4/lstmp.py b/egs/librispeech/ASR/pruned_transducer_stateless4/lstmp.py
new file mode 120000
index 000000000..9aa06f82f
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless4/lstmp.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless3/lstmp.py
\ No newline at end of file
diff --git a/egs/librispeech/ASR/pruned_transducer_stateless4/scaling_converter.py b/egs/librispeech/ASR/pruned_transducer_stateless4/scaling_converter.py
new file mode 120000
index 000000000..3b667058d
--- /dev/null
+++ b/egs/librispeech/ASR/pruned_transducer_stateless4/scaling_converter.py
@@ -0,0 +1 @@
+../pruned_transducer_stateless3/scaling_converter.py
\ No newline at end of file