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

# Install Tensorflow, check compatibility here: https://www.tensorflow.org/install/gpu 
RUN conda install --quiet --yes \
    'tensorflow-gpu=2.1*' \
    'keras-gpu' && \
    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

# Check compatibility here: https://pytorch.org/get-started/locally/
RUN conda install --quiet --yes \
     pytorch \ 
     torchvision \
     cudatoolkit=10.1 -c pytorch
#    pip install torch_nightly -f https://download.pytorch.org/whl/nightly/cu90/torch_nightly.html && \

# Clean installation
RUN conda clean --all -f -y && \
    fix-permissions $CONDA_DIR && \
    fix-permissions /home/$NB_USER