bytedeco/javacpp-presets

[PyTorch] model always loaded to GPU device 0

haifengl opened this issue · 1 comments

haifengl commented

I have a module that was moved to a GPU device (says 4). When I load the model for a file (Module.load(InputArchive), the model goes back to device 0.

haifengl commented

InputArchive.load_from(path, DeviceOptional) solves the problem.

  • 👍1