torch.Tensor.to() でGPUに変数が転送できない？model.to(device)とは返り値が違う

はじめに

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 であることが原因である．

参考文献

• torch.Tensor.to https://pytorch.org/docs/stable/generated/torch.Tensor.to.html
• torch.nn.Module.to https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.to