diff --git a/egs/libriheavy/ASR/zipformer_prompt_asr/decoder.py b/egs/libriheavy/ASR/zipformer_prompt_asr/decoder.py
index 8a55eb5c8..91f167204 100644
--- a/egs/libriheavy/ASR/zipformer_prompt_asr/decoder.py
+++ b/egs/libriheavy/ASR/zipformer_prompt_asr/decoder.py
@@ -99,6 +99,7 @@ class Decoder(nn.Module):
             # To avoid `RuntimeError: Module 'Decoder' has no attribute 'conv'`
             # when inference with torch.jit.script and context_size == 1
             self.conv = nn.Identity()
+            self.balancer2 = nn.Identity()
 
     def forward(self, y: torch.Tensor, need_pad: bool = True) -> torch.Tensor:
         """
diff --git a/egs/librispeech/ASR/zipformer/decoder.py b/egs/librispeech/ASR/zipformer/decoder.py
index 492d63fc5..7ce44495b 100644
--- a/egs/librispeech/ASR/zipformer/decoder.py
+++ b/egs/librispeech/ASR/zipformer/decoder.py
@@ -98,6 +98,7 @@ class Decoder(nn.Module):
             # To avoid `RuntimeError: Module 'Decoder' has no attribute 'conv'`
             # when inference with torch.jit.script and context_size == 1
             self.conv = nn.Identity()
+            self.balancer2 = nn.Identity()
 
     def forward(self, y: torch.Tensor, need_pad: bool = True) -> torch.Tensor:
         """