Python 后端: "预期所有张量都位于同一设备,但发现至少有两个设备,分别是 cpu 和 cuda:0"

问题描述

我在尝试使用 ROCm 为我的 AMD GPU 运行 PyTorch 时，发现了 Python 后端的这个 bug，尽管我认为 CUDA 也可能存在相同的问题。如果 PyTorch 看到 GPU，则 model.py 中的 pb2torch 会在 CPU 上创建张量，而不将其传输到 GPU。torch2pb 在转换为 numpy 之前也不会将张量从 GPU 传输回去。如果不这样做，就会出现标题中的错误。