Installing PyTorch on a Grace Hopper GH200 node with GPU support

1

Having looked at some previous threads about this, I found that when using Python 3.11, without any extra specifications to pip, the installation of PyTorch with GPU support on a GH200 node would end up with PyTorch 2.0.1 with only CPU support.

To force CUDA 12.6 and Python 3.11, I had to specify:

pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu126 --python-version 311 --only-binary=:all: --target gh200-2025-02-18-3/lib64/python3.11/site-packages/

Just putting this here in case it’s of use to others, and perhaps if I’m doing this wrongly, someone who knows better can suggest an improvement.

2

Yes, to install the ARM+CUDA binaries, you should specify the index-url pointing to the latest CUDA runtime version we support.

3

Sure - it was more the extra parameters I needed to add that I wanted to highlight, bearing in mind a previous comment of yours in a thread from last year that you were surprised pip hadn’t found the Python 3.11 wheels automatically.

4

Has PyTorch on the Grace Hopper GH200 been released? I haven’t found any other information besides this post.

5

Yes, you can install the nightly or latest stable releases for ARM + CUDA.

6

it seems pip install torch==2.6 --index-url https://download.pytorch.org/whl/cu126 works, but the wheels are not available in pypi index.