ToTensor Device

[bhashemian]

Is your feature request related to a problem? Please describe.
Currently when converting from other formats that can reside on gpu like cupy to torch tensor, using ToTensor, it will create the tensor on the cpu, and then it should be sent to gpu using ToDevice. This will create unnecessary between cpu and gpu and reduce the performance. A solution to that is to have ToTensor accept device parameter and directly create the tensor on that machine, and in this case on gpu.

[Nic-Ma]

Hi @drbeh ,

May I know when you converting a CuPy GPU array to Tensor with ToTensor, is the result a Tensor on GPU or CPU?
I suspect torch.as_tensor(CuPy GPU array) will generate a GPU Tensor.

Thanks.

[bhashemian]

Hi @Nic-Ma,

Based on my experiment, the result will reside on cpu. I made a PR (#2926) to add device to ToTensor. This also can be helpful in other situations.

Thanks

[bhashemian]

it's also easy to check

import torch
import cupy as cp
a = torch.randn(2,2).to('cuda')
b = cp.asarray(b)
c = torch.as_tensor(b)
print(a.device, '->', b.device, '->', c.device)

cuda:0 -> <CUDA Device 0> -> cpu