Bug fix in scripting mode

This commit is contained in:
Daniel Povey 2022-12-02 20:28:17 +08:00
parent 9a2a58e20d
commit 84f51ab1b1

View File

@ -1272,8 +1272,8 @@ class Swoosh(torch.nn.Module):
"""Return tan-swish activation function which is tanh(x) sigmoid(x-1)n """Return tan-swish activation function which is tanh(x) sigmoid(x-1)n
""" """
if torch.jit.is_scripting(): if torch.jit.is_scripting():
one = torch.tensor(1.0, dtype=x.dtype, device=x.device) zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
return torch.logaddexp(one, x - 1.125) - 0.08 * x - 0.3 return torch.logaddexp(zero, x - 1.125) - 0.08 * x - 0.3
return SwooshFunction.apply(x) return SwooshFunction.apply(x)