diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
index c6d44bbe1..f5bbe9d48 100644
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
@@ -978,7 +978,7 @@ class CompactRelPositionalEncoding(torch.nn.Module):
 
         """
 
-        shift = 0 if self.training else int(round(random.normalvariate(0, 1) * float(self.random_shift)))
+        shift = int(round(random.normalvariate(0, 1) * float(self.random_shift))) if self.training else 0
         self.extend_pe(x, shift)
 
         pos_emb = self.pe[