PyTorchでCUDAを使ってGPUで学習をさせたときに「Expected all tensors to be on the same device, but found at least two devices」が出たときの対処

GPUを積んだローカルPCでPyTorchの転移学習を試した。以下を参考にCUDAでGPUを使用して学習させた。

www.amazon.co.jp

結果、Optimizerの実行箇所でタイトルのエラーが発生した。
「tensor使うなら全部GPUでやらなきゃいけないのにCPUが混ざってるぞ」と言われているのは分かったけど対処法がわからない。

optimizer.to(device)

とやってみたけどダメ。 最終的に以下のページにたどり着いた。

discuss.pytorch.org

以下の関数を定義した。

def optimizer_to(optim, device):
    for param in optim.state.values():
        # Not sure there are any global tensors in the state dict
        if isinstance(param, torch.Tensor):
            param.data = param.data.to(device)
            if param._grad is not None:
                param._grad.data = param._grad.data.to(device)
        elif isinstance(param, dict):
            for subparam in param.values():
                if isinstance(subparam, torch.Tensor):
                    subparam.data = subparam.data.to(device)
                    if subparam._grad is not None:
                        subparam._grad.data = subparam._grad.data.to(device)

その後、

optimizer_to(optimizer, device)

としたことで無事学習が完了した。 ありがとう、ハンス。