【PyTorch】Expected all tensors to be on the same device（以下略）エラーの対処法

状況

PyTorch使用時に以下のようなエラーが出る。

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

対処法

スクリプトに以下を追加。

import torch
torch.set_default_tensor_type('torch.cuda.FloatTensor')

参考：https://discuss.pytorch.org/t/expected-device-cuda-0-but-got-device-cpu-though-the-two-tensors-are-on-cuda/60180