Tech

PyTorchの「Expected all tensors to be on the same device」の直し方

学習を回すと、こう怒られることがある。

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

モデルはGPUに送ったのに、入力データやラベルがCPUのまま（またはその逆）になっている、というのが原因だ。PyTorchは別々のデバイスにあるテンソル同士を計算できない。

device を1つ決めて統一する

最初に device を決めて、モデルもデータも全部そこへ送るのが基本形。

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)


for x, y in loader:
    x = x.to(device)
    y = y.to(device)
    out = model(x)
    loss = criterion(out, y)

自分で作るテンソルもデバイスを合わせる

途中で新しく作るテンソルは、CPUに生まれることがあるので明示する。

w = torch.zeros(10, device=device)      # device を指定
idx = torch.tensor([0, 1], device=device)

まとめ

  • 原因はモデルとデータのデバイス不一致
  • device を1つ決め、model.to(device) と各バッチの .to(device) を徹底
  • 途中で作るテンソルも device= で合わせる
  • 迷ったら print(x.device, next(model.parameters()).device) でどこがずれているか確認