Update Dockerfile

Dockerfile

@@ -1,70 +1,4 @@
-ARG IMAGE_NAME
-FROM ${IMAGE_NAME}:11.8.0-runtime-ubuntu20.04 as base
-
-ENV NV_CUDA_LIB_VERSION "11.8.0-1"
-
-FROM base as base-amd64
-
-ENV NV_CUDA_CUDART_DEV_VERSION 11.8.89-1
-ENV NV_NVML_DEV_VERSION 11.8.86-1
-ENV NV_LIBCUSPARSE_DEV_VERSION 11.7.5.86-1
-ENV NV_LIBNPP_DEV_VERSION 11.8.0.86-1
-ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-11-8=${NV_LIBNPP_DEV_VERSION}
-
-ENV NV_LIBCUBLAS_DEV_VERSION 11.11.3.6-1
-ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-11-8
-ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}
-
-ENV NV_NVPROF_VERSION 11.8.87-1
-ENV NV_NVPROF_DEV_PACKAGE cuda-nvprof-11-8=${NV_NVPROF_VERSION}
-
-ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
-ENV NV_LIBNCCL_DEV_PACKAGE_VERSION 2.16.2-1
-ENV NCCL_VERSION 2.16.2-1
-ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda11.8
-FROM base as base-arm64
-
-ENV NV_CUDA_CUDART_DEV_VERSION 11.8.89-1
-ENV NV_NVML_DEV_VERSION 11.8.86-1
-ENV NV_LIBCUSPARSE_DEV_VERSION 11.7.5.86-1
-ENV NV_LIBNPP_DEV_VERSION 11.8.0.86-1
-ENV NV_LIBNPP_DEV_PACKAGE libnpp-dev-11-8=${NV_LIBNPP_DEV_VERSION}
-
-ENV NV_LIBCUBLAS_DEV_PACKAGE_NAME libcublas-dev-11-8
-ENV NV_LIBCUBLAS_DEV_VERSION 11.11.3.6-1
-ENV NV_LIBCUBLAS_DEV_PACKAGE ${NV_LIBCUBLAS_DEV_PACKAGE_NAME}=${NV_LIBCUBLAS_DEV_VERSION}
-
-ENV NV_LIBNCCL_DEV_PACKAGE_NAME libnccl-dev
-ENV NV_LIBNCCL_DEV_PACKAGE_VERSION 2.16.2-1
-ENV NCCL_VERSION 2.16.2-1
-ENV NV_LIBNCCL_DEV_PACKAGE ${NV_LIBNCCL_DEV_PACKAGE_NAME}=${NV_LIBNCCL_DEV_PACKAGE_VERSION}+cuda11.8
-
-
-FROM base-${TARGETARCH}
-
-ARG TARGETARCH
-
-LABEL maintainer "NVIDIA CORPORATION <[email protected]>"
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    libtinfo5 libncursesw5 \
-    cuda-cudart-dev-11-8=${NV_CUDA_CUDART_DEV_VERSION} \
-    cuda-command-line-tools-11-8=${NV_CUDA_LIB_VERSION} \
-    cuda-minimal-build-11-8=${NV_CUDA_LIB_VERSION} \
-    cuda-libraries-dev-11-8=${NV_CUDA_LIB_VERSION} \
-    cuda-nvml-dev-11-8=${NV_NVML_DEV_VERSION} \
-    ${NV_NVPROF_DEV_PACKAGE} \
-    ${NV_LIBNPP_DEV_PACKAGE} \
-    libcusparse-dev-11-8=${NV_LIBCUSPARSE_DEV_VERSION} \
-    ${NV_LIBCUBLAS_DEV_PACKAGE} \
-    ${NV_LIBNCCL_DEV_PACKAGE} \
-    && rm -rf /var/lib/apt/lists/*
-
-# Keep apt from auto upgrading the cublas and nccl packages. See https://gitlab.com/nvidia/container-images/cuda/-/issues/88
-RUN apt-mark hold ${NV_LIBCUBLAS_DEV_PACKAGE_NAME} ${NV_LIBNCCL_DEV_PACKAGE_NAME}
-ENV LIBRARY_PATH /usr/local/cuda/lib64/stubs
-
-
+FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-runtime
 WORKDIR /app
 
 COPY ./requirements.txt /app/requirements.txt