TL;DR

PyTorchでForward/Backward処理を記述するとき、テンソルを新しく生成する場合は、計算相手と同じデバイスを設定するとよいです。

課題

PyTorchでForward/Backward処理を書いてるとき。

テンソル同士の演算で...

例

v2 =  torch.tensor(v2).clone().detach()
torch.bitwise_not(torch.bitwise_xor(v1, v2))

以下のようなエラーに遭遇することがあります。

エラー

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

演算対象の２つのテンソルv1, v2が異なるデバイス（CPUとGPU）に存在するのが原因です。

PyTorchでのdeviceの扱いは、初期化時に実行対象のデバイスを指定するのがよく見かける対応です。

実行対象デバイスを指定

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
a = torch.tensor([1., 2.]).to(device)

しかし、Forward/Backward処理はモデル定義のコードなので、記述時点ではどのデバイスで実行するのか分かりません。

そこで、Forward/Backward中に作成するテンソルは、.to(v1.device)を付けて計算相手と同じデバイスに送るようにすると良いです。

修正した例

v2 =  torch.tensor(v2).clone().detach().to(v1.device)
torch.bitwise_not(torch.bitwise_xor(v1, v2))

これで、モデル定義にデバイスを明示せずに２つのテンソルを同じデバイス上で計算することができるようになりました。