yolov5 export torchscript error: RuntimeError: Input type (torch.cuda.FloatTensor) and weight type

The yolov5 v6.0 version code of github ultralytics exported the torchscript export successfully, but when the detect.py reasoning reported RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

There are two solutions here

1. Unify the input and weight on one device (the one I used is relatively simple and rude)

Check the code and find that there is a DetectMultiBanckend in detect.py, which is a class used to load different types of models. This class is defined in the root directory /models/common.py. Check the jit (ie torchscript) corresponding to this function. place code

 It can be seen that one parameter is missing here, map_location=device, just add it

 

2. Reference (97 messages) Pytorch avoids the pit: RuntimeError: Input type(torch.cuda.FloatTensor) and weight type(torch.FloatTensor) shoul_warm boy will fly blog-CSDN blog

Guess you like

Origin blog.csdn.net/m0_52948781/article/details/124671318