Hi, I am wondering how PyTorch ensures Ops like arange or randperm which creates Tensors by taking scalars are running on GPU. Because they don’t have a context where for some op the input tensors can provide such a device context or parametric op whose parameters can provide it as well.
I’ve run the following and the output is CPU. I know that because by default the output tensor will be transferred to CPU. I want to know whether the generation process is conducted on GPU. And it seems that Ops like these are exported directly from cpp.
class _test_gpu_(torch.nn.Module):
    def __init__(self):
        super(_test_gpu_, self).__init__()
        self.f=torch.nn.Linear(4,4)
        new_Wight = torch.Tensor(np.ones([4,4])) 
        self.f.weight = torch.nn.Parameter(new_Wight)

    def forward(self,x):
        ind=torch.randperm(4)
        print(ind.device)
        y=self.f(x)   
        return  y[ind]
t= _test_gpu_().cuda()
x=torch.Tensor(np.ones([1,4])).cuda()
res=t(x)