if th.cuda.is_available(): print('yes') device = th.device("cuda") x = th.zeros([2, 3]) # cpu print(x) x = x.to(device) # GPU方法一 print(x) x = th.ones([2, 3], device=device) # GPU方法二 print(x)
from pytorch tutorial # let us run this cell only if CUDA is available # We will use ``torch.device`` objects to move tensors in and out of GPU if torch.cuda.is_available(): device = torch.device("cuda") # a CUDA device object x = torch.zeros([2,3]) # x = torch.zeros([2,3],device=device) y = torch.ones_like(x, device=device) # directly create a tensor on GPU x = x.to(device) # or just use strings ``.to("cuda")`` z = x + y print(z) print(z.to("cpu", torch.double)) # ``.to`` can also change dtype together!