14. torch cuda
device = "cuda:1" if torch.cuda.is_available() else "cpu"
x = torch.tensor([1,2,3]).to(device)
y = torch.tensor([10,20,30]).to(device)
z = x + y
print(z)
print(z.to("cpu"))
=> tensor([11, 22, 33], device='cuda:1')
=> tensor([11, 22, 33])
tensorをgpuに転送
deviceを定義することでコードが簡潔になる
15. torch cuda
device = "cuda:1" if torch.cuda.is_available() else "cpu"
x = torch.tensor([1,2,3]).to(device)
y = torch.tensor([10,20,30])
z = x + y
=> RuntimeError: expected device cuda:1 and dtype Long but got
device cpu and dtype Long
cpu + gpuの演算はエラーとなるので注意
x.to(gpu) + y.to(gpu)もしくはx.to(cpu) + y.to(cpu)