Only tensors or tuples of tensors can be output from traced functions
错误代码:
heads = {'hm': 5, 'wh': 2, 'hps': 2}
model= get_pose_net(34,heads,64)
# model = torch.load('34.weights')
# model.cuda()
model.eval()
example = torch.rand(1, 3, 512, 512)
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("cxx_model.pt")
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)