deepjavalibrary / djl

An Engine-Agnostic Deep Learning Framework in Java
https://djl.ai
Apache License 2.0
4.12k stars 654 forks source link

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

Closed javst closed 2 years ago

javst commented 2 years ago

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 commented 2 years ago

@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 commented 2 years ago

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