PyTorch on XLA Devices — PyTorch/XLA master documentationPyTorch runs on XLA devices, like TPUs, with the torch_xla package. This document describes how to run your models on these devices.https://pytorch.org/xla/release/2.1/index.html
