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