ARG PYTORCH="1.9.1"
ARG CUDA="11.1"
ARG CUDNN="8"

FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel

RUN export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 8.0+PTX" \
	&& export TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
	&& export CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" 

# INIT OF FIX FOR GPG KEY -- Hopefully they will fix this soon
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
# END OF FIX FOR GPG KEY

# apt and apt-get update
RUN apt-get update && apt update && apt upgrade -y && apt-get clean
# apt and apt-get installations
RUN apt install -y vim git unzip \
	&& apt-get install -y libosmesa6-dev build-essential libgl1-mesa-dev libglu1-mesa-dev freeglut3-dev

RUN mkdir /software
ENV SW_PATH=/software
WORKDIR $SW_PATH

ENV CONDA_DIR /opt/conda
ENV PATH $CONDA_DIR/bin:$PATH

RUN echo "Installing basic conda environment..." \
	&& conda create --name main python=3.9.5 \
	&& . /opt/conda/etc/profile.d/conda.sh \
	&& conda activate main \
	&& conda install numpy \
	&& pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html \
	&& pip install tensorflow==2.8.0 \
	&& pip install opencv-contrib-python pyyaml yacs tqdm colorama matplotlib cython tensorboardX imageio imageio-ffmpeg munkres pandas seaborn h5py sklearn