Using ATen, create CUDA tensors on specific device

1

I can use the ATen c++ library to create a tensor on the current device, e.g.

Tensor d = CUDA(kFloat).ones({3, 4})

Is it possible to create the tensor on a specific device.
i.e. to avoid calling cudaSetDevice(d), then create the tensor

Thanks

1 Like
2

I had the exact same question here. Adding new_ones to ATen would have solved this.