Creating tensors on GPU directly

1

Hi, is there a good way of constructing tensors on GPU?

Say, torch.zeros(1000, 1000).cuda() is much slower than torch.zeros(1, 1).cuda.expand(1000, 1000), but the latter is ugly.

1 Like
2 
torch.cuda.FloatTensor(1000, 1000).fill_(0)
12 Likes
3

Thanks! Could not find that in the docs. There are so many cool things in pytorch!

4

Another question: is there a cuda version of torch.arange?

7 Likes
5

Hello, but how to specify which gpu?

6

You could go with the following approach:

print("Outside device is 0")  # On device 0 (default in most scenarios)
with torch.cuda.device(1):
    print("Inside device is 1")  # On device 1
print("Outside device is still 0")  # On device 0

You could also use .cuda(GPU_NUMBER) at the end of the tensors, but the first option has been quicker and cleaner for me.

3 Likes