XLA tends tends to be better optimized for TPUs, Pytorch is better with GPUs, but I believe you can choose a backend when using Nx.