diff --git a/tools/toolchain/Dockerfile b/tools/toolchain/Dockerfile
index 306e1d3a2b..1b435f6bf4 100644
--- a/tools/toolchain/Dockerfile
+++ b/tools/toolchain/Dockerfile
@@ -11,11 +11,15 @@ ARG BASE_IMAGE
 
 USER root
 
-# install Linux packages.
+# Install Linux packages.
 COPY ./install_requirements*.sh ./
 RUN ./install_requirements.sh ${BASE_IMAGE}
 
-# copy helper scripts
+# Install Tensorflow.
+COPY ./install_tensorflow.sh ./
+RUN ./install_tensorflow.sh ${BASE_IMAGE}
+
+# Copy helper scripts.
 WORKDIR /opt/cp2k-toolchain
 RUN mkdir scripts
 COPY ./scripts/VERSION \
diff --git a/tools/toolchain/Dockerfile.cuda b/tools/toolchain/Dockerfile.cuda
index 00e90c349e..5636746f24 100644
--- a/tools/toolchain/Dockerfile.cuda
+++ b/tools/toolchain/Dockerfile.cuda
@@ -9,10 +9,14 @@ FROM nvidia/cuda:11.3.1-devel-ubuntu20.04
 ENV CUDA_PATH /usr/local/cuda
 ENV LD_LIBRARY_PATH /usr/local/cuda/lib64
 
-# install Ubuntu packages.
+# Install Ubuntu packages.
 COPY ./install_requirements_ubuntu.sh .
 RUN ./install_requirements_ubuntu.sh
 
+# Install Tensorflow.
+COPY ./install_tensorflow.sh ./
+RUN ./install_tensorflow.sh ubuntu
+
 # Install some more Ubuntu packages.
 RUN apt-get update -qq && apt-get install -qq --no-install-recommends \
     gfortran                                                          \
diff --git a/tools/toolchain/install_tensorflow.sh b/tools/toolchain/install_tensorflow.sh
new file mode 100755
index 0000000000..90ce063b1c
--- /dev/null
+++ b/tools/toolchain/install_tensorflow.sh
@@ -0,0 +1,30 @@
+#!/bin/bash -e
+
+# author: Ole Schuett
+
+if (($# != 1)); then
+  echo "Usage: install_tensorflow.sh <BASE_IMAGE>"
+  exit 1
+fi
+
+BASE_IMAGE=$1
+
+echo "Installing Tensorflow..."
+
+if [[ ${BASE_IMAGE} == ubuntu* ]]; then
+  apt-get update -qq
+  apt-get install -qq --no-install-recommends python3-pip
+  rm -rf /var/lib/apt/lists/*
+  pip install --quiet tensorflow
+
+elif [[ ${BASE_IMAGE} == fedora* ]]; then
+  dnf -qy install python3-pip
+  dnf clean -q all
+  pip install --quiet tensorflow
+
+else
+  echo "Unknown base image: ${BASE_IMAGE}"
+  exit 1
+fi
+
+#EOF