Install PyTorch for ROCm#

Refer to this section for the recommended PyTorch via PIP installation method, as well as Docker-based installation.

AMD recommends the PIP install method to create a PyTorch environment when working with ROCm™ for machine learning development.

Check for latest PIP install instructions and availability. See Compatibility matrices for support information.

The latest version of Python module numpy v2.0 is incompatible with the torch wheels for this version. Downgrade to an older version is required.
Example: pip3 install numpy==1.26.4

Install PyTorch via PIP

  1. Enter the following command to unpack and begin set up.

    sudo apt install python3-pip -y
  2. Enter this command to update the pip wheel.

    pip3 install --upgrade pip wheel
  3. Enter this command to install Torch and Torchvision for ROCm AMD GPU support.

    pip3 uninstall torch torchvision pytorch-triton-rocm numpy
    pip3 install torch-2.1.2+rocm6.1.3-cp310-cp310-linux_x86_64.whl torchvision-0.16.1+rocm6.1.3-cp310-cp310-linux_x86_64.whl pytorch_triton_rocm-2.1.0+rocm6.1.3.4d510c3a44-cp310-cp310-linux_x86_64.whl numpy==1.26.4

    This may take several minutes.

    Important! AMD recommends proceeding with ROCm WHLs available at The ROCm WHLs available at are not tested extensively by AMD as the WHLs change regularly when the nightly builds are updated.

Next, verify your PyTorch installation.

Using Docker provides portability, and access to a prebuilt Docker container that has been rigorously tested within AMD. Docker also cuts down compilation time, and should perform as expected without installation issues.

Prerequisites to install PyTorch using Docker

Docker for Ubuntu® must be installed.

To install Docker for Ubuntu, enter the following command:

sudo apt install

Use Docker image with pre-installed PyTorch

Follow these steps for installing using a Docker image.

  1. Enter the following command to pull the public PyTorch Docker image.

    sudo docker pull rocm/pytorch:rocm6.1.3_ubuntu22.04_py3.10_pytorch_release-2.1.2

    Optional: You can also download a specific and supported configuration with different user-space ROCm versions, PyTorch versions, and supported operating systems.

    Refer to to download the PyTorch Docker image.

  2. Start a Docker container using the downloaded image.

    sudo docker run -it \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      --device=/dev/kfd \
      --device=/dev/dri \
      --group-add video \
      --ipc=host \
      --shm-size 8G \

    This will automatically download the image if it does not exist on the host. You can also pass the -v argument to mount any data directories from the host onto the container.

Next, verify the PyTorch installation.

See PyTorch Installation for ROCm for more information.

Verify PyTorch installation#

Confirm if PyTorch is correctly installed.

  1. Verify if Pytorch is installed and detecting the GPU compute device.

    python3 -c 'import torch' 2> /dev/null && echo 'Success' || echo 'Failure'

    Expected result:

  2. Enter command to test if the GPU is available.

    python3 -c 'import torch; print(torch.cuda.is_available())'

    Expected result:

  3. Enter command to display installed GPU device name.

    python3 -c "import torch; print(f'device name [0]:', torch.cuda.get_device_name(0))"

    Expected result: Example: device name [0]: Radeon RX 7900 XTX

    device name [0]: <Supported AMD GPU>
  4. Enter command to display component information within the current PyTorch environment.

    python3 -m torch.utils.collect_env

    Expected result:

    PyTorch version
    ROCM used to build PyTorch
    Is CUDA available
    GPU model and configuration
    HIP runtime version
    MIOpen runtime version

Environment set-up is complete, and the system is ready for use with PyTorch to work with machine learning models, and algorithms.