LABEL maintainer="Christoph Schranz <christoph.schranz@salzburgresearch.at>"

# Install cudatoolkit must match the version of the base image
RUN conda install --quiet --yes \
    cudatoolkit=10.1 && \
    conda clean --all -f -y && \
    fix-permissions $CONDA_DIR && \
    fix-permissions /home/$NB_USER

# Install Tensorflow
RUN conda install --quiet --yes \
    'tensorflow-gpu' \
    'keras-gpu' && \
    conda clean --all -f -y && \
    fix-permissions $CONDA_DIR && \
    fix-permissions /home/$NB_USER

# Install PyTorch with dependencies
RUN conda install --quiet --yes \
    pyyaml mkl mkl-include setuptools cmake cffi typing \
    pytorch \
    torchvision && \
    conda clean --all -f -y && \
    pip install torch_nightly -f https://download.pytorch.org/whl/nightly/cu90/torch_nightly.html && \
    fix-permissions $CONDA_DIR && \
    fix-permissions /home/$NB_USER