【Pytorch Tutorial】CUDA

x = torch.cuda.FloatTensor(1) # x.get_device() == 0

y = torch.FloatTensor(1).cuda() # y.get_device() == 0

with torch.cuda.device(1): # allocates a tensor on GPU 1

    a = torch.cuda.FloatTensor(1) # transfers a tensor from CPU to GPU 1

    b = torch.FloatTensor(1).cuda() # a.get_device() == b.get_device() == 1

    c = a + b # c.get_device() == 1

    z = x + y # z.get_device() == 0

    # even within a context, you can give a GPU id to the .cuda call

    d = torch.randn(2).cuda(2) # d.get_device() == 2

 

cuda默认GPU为0，操作只能在同一个GPU上操作（默认情况下，不支持跨GPU操作）