xla Automatically move CUDA non XLA Tensors to XLA Device and back to CUDA device

Automatically move CUDA non XLA Tensors to XLA Device and back to CUDA device

Open changm opened this issue 11 months ago • 2 comments

Currently only works for inference. The assumptions don't hold for training with Autograd yet.

Feb 29 '24 00:02 changm

Is this an experimental pr or you want to merge this?

Mar 07 '24 18:03 JackCaoG

Is this an experimental pr or you want to merge this?

Ideally we would merge this, or is there a reason not to?

Mar 07 '24 22:03 changm

Sorry for being late. It's looking good!

Mar 13 '24 18:03 vanbasten23