Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

[javst]

Description

(A clear and concise description of what the bug is.) when i run the yolov5 model on my windows pc with cuda 11.2 i also use optDevice(DEvice.gpu()) but the program always tell me Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Expected Behavior

(what's the expected behavior?)

Error Message

(Paste the complete error message, including stack trace.) Loading: 100% |========================================| model devicegpu(0) 1 NDArray devicegpu(0) Exception in thread "main" ai.djl.translate.TranslateException: ai.djl.engine.EngineException: The following operation failed in the TorchScript interpreter. Traceback of TorchScript, serialized code (most recent call last): File "code/torch/models/yolo.py", line 71, in forward

[frankfliu]

@javst

You might need to set "mapLocation" option. Can you try the following:

Criteria<Image, Classifications> criteria = Criteria.builder()
      .setTypes(Image.class, Classifications.class)
      .optModelPath(Paths.get("build/pytorch_models/yolo"))
      .optOption("mapLocation", "true") 
      .optTranslator(translator)
      .optDevice(Device.gpu());

[javst]

I got it just need a optOption() thank you !!!