minor fixes for torch.jit.script

egs/librispeech/ASR/pruned_transducer_stateless/decoder.py

@@ -74,6 +74,10 @@ class Decoder(nn.Module):
                 groups=embedding_dim,
                 bias=False,
             )
+        else:
+            # To avoid `RuntimeError: Module 'Decoder' has no attribute 'conv'`
+            # when inference with torch.jit.script and context_size == 1
+            self.conv = nn.Identity()
         self.output_linear = nn.Linear(embedding_dim, vocab_size)
 
     def forward(self, y: torch.Tensor, need_pad: bool = True) -> torch.Tensor:

egs/librispeech/ASR/transducer_stateless/decoder.py

@@ -71,6 +71,10 @@ class Decoder(nn.Module):
                 groups=embedding_dim,
                 bias=False,
             )
+        else:
+            # To avoid `RuntimeError: Module 'Decoder' has no attribute 'conv'`
+            # when inference with torch.jit.script and context_size == 1
+            self.conv = nn.Identity()
 
     def forward(self, y: torch.Tensor, need_pad: bool = True) -> torch.Tensor:
         """

egs/librispeech/ASR/zipformer/decoder.py

@@ -17,7 +17,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 from scaling import Balancer
 
 
@@ -95,6 +94,10 @@ class Decoder(nn.Module):
                 max_abs=1.0,
                 prob=0.05,
             )
+        else:
+            # To avoid `RuntimeError: Module 'Decoder' has no attribute 'conv'`
+            # when inference with torch.jit.script and context_size == 1
+            self.conv = nn.Identity()
 
     def forward(self, y: torch.Tensor, need_pad: bool = True) -> torch.Tensor:
         """