non_blocking in .to() of XLA tensor

[ydcjeff]

❓ Questions and Help

I am just adapting the code that runs on cuda to xla.
Following the Running on a Single XLA Device code sinppet, and I have non_blocking=True for cuda device.
Question is can non_blocking=True convert to xla tensor asynchronously like cuda tensor does?

Thank you!

[jysohn23]

Check out how PyTorch/XLA runtime works: https://github.com/pytorch/xla/blob/master/API_GUIDE.md#xla-tensors-are-lazy

When on 'xla' device, during your model calls, instead of actually executing the ops, we build and IR graph which only gets executed once you call xm.mark_step(), which actually gets taken care of as long as you use pl.MpDeviceLoader:

xla/torch_xla/distributed/parallel_loader.py

Line 172 in 7671584

class MpDeviceLoader(object):

[ydcjeff]

Thank you @jysohn23