ARG CUDA_VERSION=10.1
ENV CUDA_VERSION=$CUDA_VERSION
RUN apt-get update && apt-get install -y --no-install-recommends \
cuda-core-$CUDA_VERSION \
cuda-cudart-$CUDA_VERSION \
&& \
rm -rf /var/lib/apt/lists/*
RUN pip install torch torchvision