zipformer2 logaddexp onnx safe (#1157)

k2-fsa · Jun 30, 2023 · 98d8946 · 98d8946
1 parent ccd8c62
commit 98d8946
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/egs/librispeech/ASR/zipformer/scaling.py b/egs/librispeech/ASR/zipformer/scaling.py
@@ -36,7 +36,9 @@ def logaddexp(x: Tensor, y: Tensor) -> Tensor:
     if not torch.jit.is_tracing():
         return torch.logaddexp(x, y)
     else:
-        return (x.exp() + y.exp()).log()
+        max_value = torch.max(x, y)
+        diff = torch.abs(x - y)
+        return max_value + torch.log1p(torch.exp(-diff))
 
 class PiecewiseLinear(object):
     """