Fix running exported model on GPU.

This commit is contained in:
Fangjun Kuang 2023-06-15 12:11:13 +08:00
parent 0ad037d076
commit 9a06745a73

View File

@ -2190,7 +2190,7 @@ class ConvolutionModule(nn.Module):
x = self.in_proj(x) # (time, batch, 2*channels)
x, s = x.chunk(2, dim=-1)
x, s = x.chunk(2, dim=2)
s = self.sigmoid(s)
x = x * s
# (time, batch, channels)