How to create tensor directly on GPU, or on device of another tensor?

Viewed 2845

I found this discussion about this, in which the code

if std.is_cuda:
    eps = torch.FloatTensor(std.size()).cuda().normal_()
else:
    eps = torch.FloatTensor(std.size()).normal_()

becomes the nice

eps = std.new().normal_()

but it is said there that

The new() method is deprecated.


  1. How to create a new tensor directly on a specific device?
  2. How to create a new tensor on the same device like another tensor without the ugly if?
2 Answers

The documentation is quite clear now about this I think. Here are described the 4 main ways to create a new tensor, and you just have to specify the device to make it on gpu :

t1 = torch.zeros((3,3), device=torch.device('cuda'))
t2 = torch.ones_like(t1, device=torch.device('cuda'))
t3 = torch.randn((3,5), device=torch.device('cuda'))

And this link adds further information about the torch.tensor() constructor. Again, the device is an argument to be specified.

If you want to use the device of another tensor, you can access it with tensor.device:

t4 = torch.empty((2,2), device=t3.device)
torch.ones_like(std).normal_()

One of the parameters of ones_like is device which (from the docs) says:

device (torch.device, optional) – the desired device of returned tensor. Default: if None, defaults to the device of input

You can also use torch.zeros_like

Related