pytorch: cuda device allocation

cuda_device_allocation.py

import torch

# check is cuda enabled
torch.cuda.is_available()

# set required device
torch.cuda.set_device(0)

# work with some required cuda device
with torch.cuda.device(1):
    # allocates a tensor on GPU 1
    a = torch.cuda.FloatTensor(1)
    assert a.get_device() == 1

    # but you still can manually assign tensor to required device
    d = torch.randn(2).cuda(2)
    assert d.get_device() == 2