Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults
错误代码:
@torch.jit.script_method
def forward(self, *x):
children = x
x = self.conv(torch.cat(x, 1))
x = self.bn(x)
if self.residual:
x += children[0]
x = self.relu(x)
return x
正确代码:
@torch.jit.script_method
def forward(self, x):
children = x
x = self.conv(torch.cat(x, 1))
x = self.bn(x)
if self.residual:
x += children[0]
x = self.relu(x)
return x