Allocate intermediate tensors in the same device as the input tensor

[polvalente]

Some functions in Torchx.Backend create intermediate Iotas and similar tensors. These should be allocated in the same device as the function's input tensor.

[josevalim]

We should probably create a helper like we did for vectorization, where do:

with_device([a, b, c], fn a, b, c ->

end)

So we can wrap our backend functions by such helper. Pinging @grzuy because he likely needs similar for candlex.