Skip to content

Instantly share code, notes, and snippets.

@ikhlestov
Created September 12, 2017 17:06
Show Gist options
  • Save ikhlestov/b62415573bfdea9bbb2aa8284f0ec805 to your computer and use it in GitHub Desktop.
Save ikhlestov/b62415573bfdea9bbb2aa8284f0ec805 to your computer and use it in GitHub Desktop.
pytorch: cuda device allocation
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment