[TPU] Correctly profile peak memory usage & Upgrade PyTorch XLA (#9438)

This commit is contained in:
Woosuk Kwon
2024-10-30 02:41:38 -07:00
committed by GitHub
parent 6aa6020f9b
commit 211fe91aa8
3 changed files with 11 additions and 10 deletions

View File

@ -56,8 +56,8 @@ First, install the dependencies:
$ pip uninstall torch torch-xla -y
$ # Install PyTorch and PyTorch XLA.
$ export DATE="20240828"
$ export TORCH_VERSION="2.5.0"
$ export DATE="20241017"
$ export TORCH_VERSION="2.6.0"
$ pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl
$ pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl