torch.jit保存，加载模型

ori_backbone.eval()
x = torch.randn(1,3,112,112).cuda()
traced_cell = torch.jit.trace(ori_backbone, (x))
torch.jit.save(traced_cell, "jit_checkpoint.pth")

loc=torch.device('cpu')
loc=torch.device('cuda:0')

model=torch.jit.load(spath[0], map_location=loc)
model.eval()

导出jit报错：

Double did not match Long

原因：可能是jit不支持条件判断。

解决方法：

self.model.model[-1].export=True

self.model.eval()
self.model=self.model.float()

self.model.model[-1].export=True
x = torch.randn(1, 3, 544, 960).float().cuda()
traced_cell = torch.jit.trace(self.model, (x))
torch.jit.save(traced_cell, "jit_face_detect.pth")

但是yolov5的不一样，导出代码如下：

"""Exports a YOLOv5 *.pt