はじめに

NNモデルをGPUに転送する際には model.to(device) で転送することができます．
一方，単なる配列を転送する際にはこれではうまくいかなかったので，備忘録として記録します．

環境

Python 3.11
torch 2.3.0+cu118

方法

これだとうまくいかない．

tensor = torch.rand([2,3])
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print(tensor.device)
# cpu

tensor.to(device)
print(tensor.device)
# cpu

model.to(divece) と違い，配列をGPUに転送する際には返り値をもらってくる必要がある．

tensor_gpu = tensor.to(device)
print(tensor_gpu.device)
# cuda:0

torch.Tensor.to の返り値は Tensor であるのに対し，torch.nn.module.to の返り値は Self であることが原因である．

参考文献