Dockerfile

ARG FROM_IMAGE=lcskrishna/rocm-pytorch:rocm3.3_ubuntu16.04_py3.6_pytorch_updated

FROM ${FROM_IMAGE}
RUN \
    git clone --recursive https://github.com/ROCmSoftwarePlatform/apex.git && \
    cd apex && \
    python3.6 setup.py install --cpp_ext --cuda_ext