[BUG]: Arm64 users need to install CUDA enabled PyTorch by hand

[dagardner-nv]

Version

25.02

Which installation method(s) does this occur on?

Docker

Describe the bug.

Pre-built wheels for torch with Cuda support for ARM are new:
https://discuss.pytorch.org/t/pytorch-arm-cuda-support/208857/4

The CUDA builds for x86_64 contain a version postfix of +cu124 in the metadata of the wheel, allowing the user to specify the CUDA build.

The wheels for ARM lack this, requiring torch to be installed with:

pip install --force-reinstall --index-url https://download.pytorch.org/whl/cu124 torch==2.4.0

However specifying --index-url in a requirements file would prevent other dependencies from being installed, and performing pip install --extra-index-url https://download.pytorch.org/whl/cu124 torch==2.4.0 will install a CPU-only version of torch.

Minimum reproducible example

pip install --extra-index-url https://download.pytorch.org/whl/cu124 torch==2.4.0+cu124

Relevant log output

Click here to see error details

[Paste the error here, it will be hidden by default]

Full env printout

Click here to see environment details

[Paste the results of print_env.sh here, it will be hidden by default]

Other/Misc.

No response

Code of Conduct

• I agree to follow Morpheus' Code of Conduct
• I have searched the open bugs and have found no duplicates for this bug report

[dagardner-nv]

The issue will be fixed for PyTorch 2.6