diff --git a/.build/Dockerfile b/.build/Dockerfile
index 2c543e3..921d119 100755
--- a/.build/Dockerfile
+++ b/.build/Dockerfile
@@ -359,28 +359,26 @@ RUN julia -e 'import Pkg; Pkg.update()' && \
 
 LABEL maintainer="Christoph Schranz <christoph.schranz@salzburgresearch.at>"
 
-# Install cudatoolkit must match the version of the base image
+# Install Tensorflow, check compatibility here: https://www.tensorflow.org/install/gpu 
 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' \
+    'tensorflow-gpu=2.1*' \
     '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 && \
+    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
 
diff --git a/src/Dockerfile.gpulibs b/src/Dockerfile.gpulibs
index 4c8e3cc..088a9bd 100644
--- a/src/Dockerfile.gpulibs
+++ b/src/Dockerfile.gpulibs
@@ -1,26 +1,24 @@
 LABEL maintainer="Christoph Schranz <christoph.schranz@salzburgresearch.at>"
 
-# Install cudatoolkit must match the version of the base image
+# Install Tensorflow, check compatibility here: https://www.tensorflow.org/install/gpu 
 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' \
+    'tensorflow-gpu=2.1*' \
     '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 && \
+    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