ROCm slim

rocm.Dockerfile

ARG UBUNTU_VER=22.04
ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/rocm5.6 
ARG ROCM_VER=5.7.1
FROM rocm/dev-ubuntu-22.04:${ROCM_VER}-complete as rocm_binary

FROM ubuntu:${UBUNTU_VER} as pytorch_builder
ARG ROCM_VER
ARG TORCH_INDEX_URL
ENV DEBIAN_FRONTEND="noninteractive"
COPY --from=rocm_binary /opt/rocm-${ROCM_VER} /opt/rocm-${ROCM_VER}

RUN apt update && \
    apt install libjpeg-dev python3-dev python3-pip -y && \
    pip install wheel setuptools && \
    pip install --no-cache-dir --pre torch torchvision torchaudio pytorch-triton-rocm --extra-index-url ${TORCH_INDEX_URL}&& \
    apt clean && rm -rf /var/lib/apt/lists/* && pip cache purge
RUN groupadd -g 109 render