rocm/pytorch:latest broken: torchvision import crashes with 'operator torchvision::nms does not exist'

[sunway513]

Summary

rocm/pytorch:latest Docker image ships a broken torchvision — import torchvision crashes immediately.

Reproducer

docker run --rm rocm/pytorch:latest python3 -c "import torchvision"

Error

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/opt/venv/lib/python3.12/site-packages/torchvision/__init__.py", line 10, in <module>
    from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils
  File "/opt/venv/lib/python3.12/site-packages/torchvision/_meta_registrations.py", line 163, in <module>
    @torch.library.register_fake("torchvision::nms")
  File "/opt/venv/lib/python3.12/site-packages/torch/library.py", line 1062, in register
    use_lib._register_fake(
  File "/opt/venv/lib/python3.12/site-packages/torch/library.py", line 210, in _register_fake
    handle = entry.fake_impl.register(
  File "/opt/venv/lib/python3.12/site-packages/torch/_library/fake_impl.py", line 50, in register
    if torch._C._dispatch_has_kernel_for_dispatch_key(self.qualname, "Meta"):
RuntimeError: operator torchvision::nms does not exist

Versions in the image

torch:        2.9.1+rocm7.2.1.lw.gitff65f5bc
torchvision:  0.24.1
Python:       3.12

Impact

• Breaks all downstream CI pipelines that depend on rocm/pytorch:latest (ATOM, vLLM OOT tests, etc.)
• Any code path that triggers import torchvision (including HuggingFace transformers image processing) will crash
• ATOM CI is currently blocked by this: https://github.com/ROCm/ATOM/actions/runs/23675231613/job/68977006040

Root cause

torchvision's C++ extension (_C.so) does not have the custom ops (torchvision::nms, torchvision::roi_align, etc.) registered. This suggests torchvision was installed from a pre-built CPU-only wheel rather than being compiled from source against the ROCm torch build.

Suggested fix

Rebuild torchvision from source against the torch in the container, or install from the ROCm-compatible wheel index:

pip install --force-reinstall torchvision --index-url https://download.pytorch.org/whl/rocm7.2

[sunway513]

Root cause identified

CUDA wheels were installed into the ROCm container.

torch:        2.9.1+rocm7.2.1  (ROCm — correct)
torchvision:  0.24.1+cu128     (CUDA — wrong!)
torchaudio:   2.9.1+cu128      (CUDA — wrong!)

Evidence:

$ docker run --rm rocm/pytorch:latest cat /opt/venv/lib/python3.12/site-packages/torchvision/version.py
__version__ = '0.24.1+cu128'

$ docker run --rm rocm/pytorch:latest ldd /opt/venv/lib/python3.12/site-packages/torchvision/_C.so | grep cuda
	libcudart.so.12 => not found
	libc10_cuda.so => not found
	libtorch_cuda.so => not found

torchvision _C.so links against CUDA libraries (libcudart.so.12, libc10_cuda.so, libtorch_cuda.so) which don't exist on the ROCm container. The custom ops (torchvision::nms, etc.) fail to register because the CUDA _C.so can't load, so import torchvision crashes.

Same issue for torchaudio (2.9.1+cu128).

Suggested fix

The Dockerfile (or build script) that produces rocm/pytorch:latest needs to install torchvision and torchaudio from a ROCm-compatible source. Either:

Option A — Install from ROCm wheel index:

pip install torchvision torchaudio --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/

Option B — Build from source:

pip install --no-build-isolation torchvision torchaudio

Option C — Use PyTorch ROCm nightly index:

pip install torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2

The current build likely does pip install torchvision torchaudio without specifying --index-url, which defaults to PyPI and pulls the CUDA +cu128 wheels.

[sunway513]

Fix verified

Important: --index-url https://download.pytorch.org/whl/rocm7.2 upgrades torch to 2.11.0 — this is NOT acceptable as a fix. It would break all downstream dependencies pinned to torch 2.9.1.

The correct fix is to use the ROCm 7.2.1 wheel index with --no-deps to replace only torchvision and torchaudio without touching torch:

pip install --force-reinstall --no-deps   --find-links https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/   "torchvision==0.24.0+rocm7.2.1.gitb919bd0c"   "torchaudio==2.9.0+rocm7.2.1.gite3c6ee2b"

Verified result

torch:        2.9.1+rocm7.2.1              (unchanged)
torchvision:  0.24.0+rocm7.2.1.gitb919bd0c (fixed)
torchaudio:   2.9.0+rocm7.2.1.gite3c6ee2b  (fixed)
→ import torchvision OK
→ InterpolationMode.BILINEAR OK

For the Dockerfile that builds rocm/pytorch:latest

The root cause is likely a pip install torchvision torchaudio without specifying the ROCm wheel source. The fix in the Dockerfile should be:

RUN pip install --no-deps   --find-links https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1/   torchvision torchaudio

Or pin explicitly to avoid future breakage.