import torch
import numpy as np
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
x = torch.tensor(np.arange(15).reshape(3,5))
if torch.cuda.is_available():
device = torch.device("cuda")
y = torch.ones_like(x,device=device)
x = x.to(device)
z = x + y
print(z)
print(z.to("cpu",torch.double))
tensor([[ 1, 2, 3, 4, 5],
[ 6, 7, 8, 9, 10],
[11, 12, 13, 14, 15]], device='cuda:0', dtype=torch.int32)
tensor([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.],
[11., 12., 13., 14., 15.]], dtype=torch.float64)