diff --git a/.github/actions/build-image/action.yml b/.github/actions/build-image/action.yml
new file mode 100644
index 0000000..efa0aba
--- /dev/null
+++ b/.github/actions/build-image/action.yml
@@ -0,0 +1,48 @@
+inputs:
+  dockerhub_username:
+    description: "The DockerHub username."
+    required: true
+
+  dockerhub_token:
+    description: "The DockerHub login token."
+    required: true
+
+  docker_file:
+    description: "The docker file."
+    required: true
+
+  tags:
+    description: "Image tags (csv)."
+    required: true
+
+  platforms:
+    description: "Platforms to build for (csv)."
+    default: "linux/amd64,linux/arm64"
+
+  push_image:
+    description: "Whether to push to DockerHub."
+    default: false
+
+runs:
+  using: "composite"
+  steps:
+    - name: Set up QEMU 🌈
+      uses: docker/setup-qemu-action@v3
+
+    - name: Set up Docker Buildx ✨
+      uses: docker/setup-buildx-action@v3
+
+    - name: Login to Docker Hub 🎪
+      uses: docker/login-action@v3
+      with:
+        username: ${{ inputs.dockerhub_username }}
+        password: ${{ inputs.dockerhub_token }}
+
+    - name: Build and deploy image 🐳
+      uses: docker/build-push-action@v5
+      with:
+        context: docker
+        file: ${{ inputs.docker_file }}
+        platforms: ${{ inputs.platforms }}
+        push: ${{ inputs.push_image }}
+        tags: ${{ inputs.tags }}
diff --git a/.github/workflows/deploy_images.yml b/.github/workflows/deploy_images.yml
index e7a334b..bdc109a 100644
--- a/.github/workflows/deploy_images.yml
+++ b/.github/workflows/deploy_images.yml
@@ -3,167 +3,185 @@ name: Deploy images
 on:
   workflow_dispatch:
     inputs:
-      skip_push:
-        description: Skip pushing images? (true|false)
-        required: false
-        default: 'false'
-      skip_docs:
-        description: Skip the docs image? (true|false)
-        required: false
-        default: 'false'
+      push_image:
+        description: Push to DockerHub
+        type: boolean
+        default: false
+
+      specific_job:
+        description: Specific job to run
+        type: choice
+        default: all
+        options:
+          - all
+          - py37
+          - py38
+          - py39
+          - py310
+          - py311
+          - py311-cuda
+          - example
+
+      platforms:
+        description: Platforms to build for
+        type: choice
+        default: linux/amd64,linux/arm64
+        options:
+          - linux/amd64,linux/arm64
+          - linux/amd64
+          - linux/arm64
 
 jobs:
   py37:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py37') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_37
-          repository: cmsml/cmsml
-          tags: "3.7"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_37
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.7
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py38:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py38') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_38
-          repository: cmsml/cmsml
-          tags: "3.8"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_38
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.8
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py39:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py39') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_39
-          repository: cmsml/cmsml
-          tags: 3.9,3,latest
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_39
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.9,cmsml/cmsml:3,cmsml/cmsml:latest
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py39_base:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py39') }}
     runs-on: ubuntu-latest
-    if: ${{ github.event.inputs.skip_push != 'true' }}
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_39_base
-          repository: cmsml/cmsml
-          tags: 3.9_base
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_39_base
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.9_base
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py310:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py310') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_310
-          repository: cmsml/cmsml
-          tags: "3.10"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_310
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.10
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
 
   py311:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py311') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_311
-          repository: cmsml/cmsml
-          tags: "3.11"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-          
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_311
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:3.11
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
+
   py311-cuda:
+    if: ${{ github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'py311-cuda') }}
     runs-on: ubuntu-latest
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_311_cuda
-          repository: cmsml/cmsml
-          tags: "3.11-cuda"
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-          
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_311_cuda
+          # gpu / cuda only available on x86, not on arm
+          platforms: linux/amd64
+          tags: cmsml/cmsml:3.11-cuda
+          push_image: ${{ github.event.inputs.push_image == 'true' }}
+
   docs:
+    if: ${{ github.event.inputs.push_image == 'true' && (github.event.inputs.specific_job == 'all' || contains(github.event.inputs.specific_job, 'example')) }}
     needs: py39
     runs-on: ubuntu-latest
-    if: ${{ github.event.inputs.skip_docs != 'true' && github.event.inputs.skip_push != 'true' }}
     steps:
       - name: Checkout 🛎️
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           persist-credentials: false
 
       - name: Build and deploy image 🐳
-        uses: docker/build-push-action@v1
-        with:
-          path: docker
-          dockerfile: docker/Dockerfile_docs
-          repository: cmsml/cmsml
-          tags: docs
-          push: ${{ github.event.inputs.skip_push != 'true' }}
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
+        uses: ./.github/actions/build-image
+        with:
+          dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+          dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+          docker_file: docker/Dockerfile_docs
+          platforms: ${{ github.event.inputs.platforms }}
+          tags: cmsml/cmsml:docs
+          push_image: true
diff --git a/.github/workflows/lint_and_test.yml b/.github/workflows/lint_and_test.yml
index 5ec6aaa..30f2ce0 100644
--- a/.github/workflows/lint_and_test.yml
+++ b/.github/workflows/lint_and_test.yml
@@ -37,6 +37,7 @@ jobs:
           - {tag: "3.9_base", tf: "2.11.1"}
           - {tag: "3.9_base", tf: "2.12.1"}
           - {tag: "3.9_base", tf: "2.13.0"}
+          - {tag: "3.9_base", tf: "2.16.1"}
     name: test (image=${{ matrix.versions.tag }}, tf=${{ matrix.versions.tf }})
     steps:
       - name: Checkout 🛎️
diff --git a/.gitignore b/.gitignore
index 5fe43c3..2dbcbc1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -19,3 +19,4 @@ build
 tmp
 docs/_build
 .ipynb_checkpoints
+.python-version
diff --git a/README.md b/README.md
index 266a272..000cacb 100644
--- a/README.md
+++ b/README.md
@@ -40,7 +40,7 @@ The documentation of this Python package is hosted on [readthedocs](http://cmsml
 **However**, note that this documentation only covers the API and technical aspects of the package itself.
 Usage examples and further techniques for working with machine learning tools in CMS, alongside a collection of useful guidelines can be found in the [general CMS ML group documentation](https://cms-ml.github.io/documentation).
 
-Click [here](https://github.com/cms-ml/cmsml/issues/new?labels=suggestion&template=feature-suggestion.md&) to submit a feature suggestion!
+Click [here](https://github.com/cms-ml/cmsml/issues/new?labels=suggestion&template=feature-suggestion.md) to submit a feature suggestion!
 
 
 <!-- marker-after-header -->
@@ -53,6 +53,15 @@ Click [here](https://github.com/cms-ml/cmsml/issues/new?labels=suggestion&templa
 
 To use the cmsml package via docker, checkout our [DockerHub](https://hub.docker.com/repository/docker/cmsml/cmsml) which contains tags for several Python versions.
 
+| Image                                                        | Python version |    TF Version    | PyTorch Version | GPU support |
+| :----------------------------------------------------------- | :------------: | :--------------: | :-------------: | :---------: |
+| `cmsml/cmsml:3.7`                                            |      3.7       |      2.11.1      |     1.13.1      |      ✘      |
+| `cmsml/cmsml:3.8`                                            |      3.8       |      2.13.1      | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.9`<br>`cmsml/cmsml:3`<br>`cmsml/cmsml:latest` |      3.9       | latest (~2.16.1) | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.10`                                           |      3.10      | latest (~2.16.1) | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.11`                                           |      3.11      | latest (~2.16.1) | latest (~2.3.0) |      ✘      |
+| `cmsml/cmsml:3.11-cuda`                                      |      3.11      | latest (~2.16.1) | latest (~2.3.0) |      ✔︎      |
+
 <!-- marker-after-docker -->
 
 <!-- marker-before-testing -->
diff --git a/cmsml/tensorflow/__init__.py b/cmsml/tensorflow/__init__.py
index a6fc657..c6be046 100644
--- a/cmsml/tensorflow/__init__.py
+++ b/cmsml/tensorflow/__init__.py
@@ -6,16 +6,16 @@
 """
 
 __all__ = [
-    "import_tf", "save_frozen_graph", "save_graph", "load_frozen_graph", "load_graph",
-    "write_graph_summary", "load_model", "load_graph_def",
+    "import_tf", "tf_version_check", "tf_keras_version_check", "save_frozen_graph", "save_graph",
+    "load_frozen_graph", "load_graph", "write_graph_summary", "load_model", "load_graph_def",
     "OpsData", "get_graph_ops",
 ]
 
 
 # provisioning imports
 from cmsml.tensorflow.tools import (
-    import_tf, save_frozen_graph, save_graph, load_frozen_graph, load_graph, write_graph_summary,
-    load_model, load_graph_def,
+    import_tf, tf_version_check, tf_keras_version_check, save_frozen_graph, save_graph,
+    load_frozen_graph, load_graph, write_graph_summary, load_model, load_graph_def,
 )
 
 from cmsml.tensorflow.aot import (
diff --git a/cmsml/tensorflow/tools.py b/cmsml/tensorflow/tools.py
index 08b08a3..20ea707 100644
--- a/cmsml/tensorflow/tools.py
+++ b/cmsml/tensorflow/tools.py
@@ -14,7 +14,7 @@
 from typing import Any
 from tensorflow.core.framework.graph_pb2 import GraphDef
 
-from cmsml.util import MockModule
+from cmsml.util import MockModule, _op_map
 
 
 tf = MockModule("tensorflow")
@@ -71,6 +71,59 @@ def import_tf(
     return tf, tf1, tf_version
 
 
+def tf_version_check(op: str, version: int | tuple[int, ...]) -> bool:
+    """
+    Compares the installed TensorFlow version with *version* using an operator *op*, which should be
+    any of ``"=="``, ``"!="``, ``"<"``, ``"<="``, ``">"`` or ``">="``. Examples:
+
+    .. code-block:: python
+
+        # actual version is 2.16
+        tf_version_check("==", (2, 16))  # -> True
+        tf_version_check(">", (2, 15))   # -> True
+        tf_version_check("!=", (2, 16))  # -> False
+    """
+    if op not in _op_map:
+        raise ValueError(f"unsupported operator '{op}'")
+    if not isinstance(version, tuple):
+        version = (version,)
+    if len(version) > 3:
+        raise ValueError("version must be at most a 3-tuple")
+    # get the tf version
+    tf_version = import_tf()[2][:len(version)]
+    # comparison
+    return _op_map[op](tf_version, version)
+
+
+def tf_keras_version_check(op: str, version: int | tuple[int, ...]) -> bool:
+    """
+    Compares the installed Keras version shipped with TensorFlow with *version* using an operator
+    *op*, which should be any of ``"=="``, ``"!="``, ``"<"``, ``"<="``, ``">"`` or ``">="``.
+    Examples:
+
+    .. code-block:: python
+
+        # actual version is 3.3
+        keras_version_check("==", (3, 3))  # -> True
+        keras_version_check(">", (3, 2))   # -> True
+        keras_version_check("!=", (3, 3))  # -> False
+    """
+    if op not in _op_map:
+        raise ValueError(f"unsupported operator '{op}'")
+    if not isinstance(version, tuple):
+        version = (version,)
+    if len(version) > 3:
+        raise ValueError("version must be at most a 3-tuple")
+    # get the keras version, which is tf version dependent
+    if tf_version_check(">=", (2, 16)):
+        keras_version = tuple(map(int, import_tf()[0].keras.__version__.split(".", 2)))
+    else:
+        import keras
+        keras_version = tuple(map(int, keras.__version__.split(".", 2)))
+    # comparison
+    return _op_map[op](keras_version, version)
+
+
 def save_graph(
     path: str,
     obj: Any,
@@ -138,16 +191,20 @@ def save_frozen_graph(
         from tensorflow.python.eager.function import ConcreteFunction
 
         if isinstance(obj, tf.keras.Model):
-            learning_phase_orig = tf.keras.backend.get_value(tf.keras.backend.learning_phase())
-            tf.keras.backend.set_learning_phase(False)
-            model_func = saving_utils.trace_model_call(obj)
-            if model_func.function_spec.arg_names and not model_func.input_signature:
-                raise ValueError(
-                    "when obj is a keras model callable accepting arguments, its "
-                    "input signature must be frozen by building the model",
-                )
-            obj = model_func.get_concrete_function()
-            tf.keras.backend.set_learning_phase(learning_phase_orig)
+            if tf_keras_version_check(">=", 3):
+                obj = obj._default_save_signature.get_concrete_function()
+            else:
+                # set the learning phase to False for the duration of the export for keras <= 3
+                learning_phase_orig = tf.keras.backend.get_value(tf.keras.backend.learning_phase())
+                tf.keras.backend.set_learning_phase(False)
+                model_func = saving_utils.trace_model_call(obj)
+                if model_func.function_spec.arg_names and not model_func.input_signature:
+                    raise ValueError(
+                        "when obj is a keras model callable accepting arguments, its "
+                        "input signature must be frozen by building the model",
+                    )
+                obj = model_func.get_concrete_function()
+                tf.keras.backend.set_learning_phase(learning_phase_orig)
 
         elif isinstance(obj, Function):
             if obj.function_spec.arg_names and not obj.input_signature:
@@ -295,8 +352,8 @@ def load_graph_def(
     serving_key: str = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
 ) -> GraphDef:
     """
-    Loads the model saved at *model_path* and returns the GraphDef of it. Supported input types are tensorflow and keras
-    SavedModels, as well as frozen graphs.
+    Loads the model saved at *model_path* and returns the GraphDef of it. Supported input types are
+    tensorflow and keras SavedModels, as well as frozen graphs.
     """
     tf, tf1, tf_version = import_tf()
 
@@ -304,8 +361,6 @@ def load_graph_def(
 
     # if model_path is directory try load as saved model
     if os.path.isdir(model_path) and tf.saved_model.contains_saved_model(model_path):
-        # if keras model try to load as keras model
-        # else load as tensorflow saved model
         loaded_saved_model = load_model(model_path)
 
         # extract graph
@@ -325,6 +380,17 @@ def load_graph_def(
 
         return graph_def
 
+    # load from new keras model interface
+    if (
+        tf_keras_version_check(">=", 3) and
+        os.path.isfile(model_path) and
+        model_path.endswith((".keras", ".h5"))
+    ):
+        loaded_keras_model = load_model(model_path)
+        # TODO: no interace yet to choose the serving key
+        concrete_function = loaded_keras_model._default_save_signature.get_concrete_function()
+        return concrete_function.graph.as_graph_def()
+
     raise FileNotFoundError(f"{model_path} contains neither frozen graph nor SavedModel")
 
 
@@ -337,7 +403,15 @@ def load_model(model_path: str) -> tf.Model:
 
     model_path = os.path.expandvars(os.path.expanduser(str(model_path)))
 
-    if os.path.isdir(model_path) and os.path.exists(os.path.join(model_path, "keras_metadata.pb")):
+    if (
+        (
+            os.path.isdir(model_path) and
+            os.path.exists(os.path.join(model_path, "keras_metadata.pb"))
+        ) or (
+            os.path.isfile(model_path) and
+            model_path.endswith((".keras", ".h5"))
+        )
+    ):
         model = tf.keras.models.load_model(model_path)
     else:
         model = tf.saved_model.load(model_path)
@@ -359,7 +433,6 @@ def write_graph_summary(
     .. note::
         When used with TensorFlow v1, eager mode must be disabled.
     """
-
     # prepare the summary dir
     if not os.path.exists(summary_dir):
         os.makedirs(summary_dir)
diff --git a/cmsml/util.py b/cmsml/util.py
index e152834..f8e7ac0 100644
--- a/cmsml/util.py
+++ b/cmsml/util.py
@@ -16,6 +16,7 @@
 import tempfile
 import contextlib
 import subprocess
+import operator
 import signal
 import importlib
 import six
@@ -33,6 +34,15 @@
     enumerate,
 )
 
+_op_map = {
+    "==": operator.eq,
+    "!=": operator.ne,
+    "<": operator.lt,
+    "<=": operator.le,
+    ">": operator.gt,
+    ">=": operator.ge,
+}
+
 
 def is_lazy_iterable(obj: Any) -> bool:
     """
diff --git a/docker/Dockerfile_310 b/docker/Dockerfile_310
index 541be8f..3fb5112 100644
--- a/docker/Dockerfile_310
+++ b/docker/Dockerfile_310
@@ -5,27 +5,26 @@ WORKDIR /root
 
 # minimal software stack
 RUN apt-get update; apt-get clean
-RUN apt-get install -y nano less htop git; apt-get clean
+RUN apt-get install -y nano less htop git libhdf5-serial-dev; apt-get clean
 
 # python software stack
 RUN pip install --no-cache-dir --upgrade pip setuptools
 RUN pip install --no-cache-dir --upgrade ipython
-RUN pip install --no-cache-dir \
-    numpy \
-    scipy \
-    matplotlib \
-    pandas \
-    numexpr \
-    jupyterlab \
-    notebook \
-    scikit-learn \
-    scikit-optimize \
-    tensorflow \
-    xgboost \
-    scinum \
-    nvidia_smi \
-    py3nvml \
-    torch
+RUN pip install --no-cache-dir numpy
+RUN pip install --no-cache-dir scipy
+RUN pip install --no-cache-dir matplotlib
+RUN pip install --no-cache-dir pandas
+RUN pip install --no-cache-dir numexpr
+RUN pip install --no-cache-dir jupyterlab
+RUN pip install --no-cache-dir notebook
+RUN pip install --no-cache-dir scikit-learn
+RUN pip install --no-cache-dir scikit-optimize
+RUN pip install --no-cache-dir tensorflow
+RUN pip install --no-cache-dir xgboost
+RUN pip install --no-cache-dir scinum
+RUN pip install --no-cache-dir nvidia_smi
+RUN pip install --no-cache-dir py3nvml
+RUN pip install --no-cache-dir torch
 
 # install cmsml from master
 RUN git clone https://github.com/cms-ml/cmsml.git && \
diff --git a/docker/Dockerfile_311 b/docker/Dockerfile_311
index 7d80cfe..64b8ca7 100644
--- a/docker/Dockerfile_311
+++ b/docker/Dockerfile_311
@@ -5,27 +5,26 @@ WORKDIR /root
 
 # minimal software stack
 RUN apt-get update; apt-get clean
-RUN apt-get install -y nano less htop git; apt-get clean
+RUN apt-get install -y nano less htop git libhdf5-serial-dev; apt-get clean
 
 # python software stack
 RUN pip install --no-cache-dir --upgrade pip setuptools
 RUN pip install --no-cache-dir --upgrade ipython
-RUN pip install --no-cache-dir \
-    numpy \
-    scipy \
-    matplotlib \
-    pandas \
-    numexpr \
-    jupyterlab \
-    notebook \
-    scikit-learn \
-    scikit-optimize \
-    tensorflow \
-    xgboost \
-    scinum \
-    nvidia_smi \
-    py3nvml \  
-    torch
+RUN pip install --no-cache-dir numpy
+RUN pip install --no-cache-dir scipy
+RUN pip install --no-cache-dir matplotlib
+RUN pip install --no-cache-dir pandas
+RUN pip install --no-cache-dir numexpr
+RUN pip install --no-cache-dir jupyterlab
+RUN pip install --no-cache-dir notebook
+RUN pip install --no-cache-dir scikit-learn
+RUN pip install --no-cache-dir scikit-optimize
+RUN pip install --no-cache-dir tensorflow
+RUN pip install --no-cache-dir xgboost
+RUN pip install --no-cache-dir scinum
+RUN pip install --no-cache-dir nvidia_smi
+RUN pip install --no-cache-dir py3nvml
+RUN pip install --no-cache-dir torch
 
 # install cmsml from master
 RUN git clone https://github.com/cms-ml/cmsml.git && \
diff --git a/docker/Dockerfile_311_cuda b/docker/Dockerfile_311_cuda
index 8315a38..5988d01 100644
--- a/docker/Dockerfile_311_cuda
+++ b/docker/Dockerfile_311_cuda
@@ -1,4 +1,4 @@
-FROM python:3.11
+FROM tensorflow/tensorflow:2.16.1-gpu
 
 # set the workdir
 WORKDIR /root
@@ -20,7 +20,6 @@ RUN pip install --no-cache-dir \
     notebook \
     scikit-learn \
     scikit-optimize \
-    tensorflow[and-cuda] \
     xgboost \
     scinum \
     nvidia_smi \
diff --git a/docker/Dockerfile_37 b/docker/Dockerfile_37
index 83fc1db..fa6de36 100644
--- a/docker/Dockerfile_37
+++ b/docker/Dockerfile_37
@@ -5,27 +5,26 @@ WORKDIR /root
 
 # minimal software stack
 RUN apt-get update; apt-get clean
-RUN apt-get install -y nano less htop git; apt-get clean
+RUN apt-get install -y nano less htop git libhdf5-serial-dev; apt-get clean
 
 # python software stack
 RUN pip install --no-cache-dir --upgrade pip setuptools
 RUN pip install --no-cache-dir --upgrade ipython
-RUN pip install --no-cache-dir \
-    numpy \
-    scipy \
-    matplotlib \
-    pandas \
-    numexpr \
-    jupyterlab \
-    notebook \
-    scikit-learn \
-    scikit-optimize \
-    tensorflow \
-    xgboost \
-    scinum \
-    nvidia_smi \
-    py3nvml \
-    torch
+RUN pip install --no-cache-dir numpy
+RUN pip install --no-cache-dir scipy
+RUN pip install --no-cache-dir matplotlib
+RUN pip install --no-cache-dir pandas
+RUN pip install --no-cache-dir numexpr
+RUN pip install --no-cache-dir jupyterlab
+RUN pip install --no-cache-dir notebook
+RUN pip install --no-cache-dir scikit-learn
+RUN pip install --no-cache-dir scikit-optimize
+RUN pip install --no-cache-dir tensorflow
+RUN pip install --no-cache-dir xgboost
+RUN pip install --no-cache-dir scinum
+RUN pip install --no-cache-dir nvidia_smi
+RUN pip install --no-cache-dir py3nvml
+RUN pip install --no-cache-dir torch
 
 # install cmsml from master
 RUN git clone https://github.com/cms-ml/cmsml.git && \
diff --git a/docker/Dockerfile_38 b/docker/Dockerfile_38
index 54cc560..e1e96db 100644
--- a/docker/Dockerfile_38
+++ b/docker/Dockerfile_38
@@ -5,27 +5,26 @@ WORKDIR /root
 
 # minimal software stack
 RUN apt-get update; apt-get clean
-RUN apt-get install -y nano less htop git; apt-get clean
+RUN apt-get install -y nano less htop git libhdf5-serial-dev; apt-get clean
 
 # python software stack
 RUN pip install --no-cache-dir --upgrade pip setuptools
 RUN pip install --no-cache-dir --upgrade ipython
-RUN pip install --no-cache-dir \
-    numpy \
-    scipy \
-    matplotlib \
-    pandas \
-    numexpr \
-    jupyterlab \
-    notebook \
-    scikit-learn \
-    scikit-optimize \
-    tensorflow \
-    xgboost \
-    scinum \
-    nvidia_smi \
-    py3nvml \
-    torch
+RUN pip install --no-cache-dir numpy
+RUN pip install --no-cache-dir scipy
+RUN pip install --no-cache-dir matplotlib
+RUN pip install --no-cache-dir pandas
+RUN pip install --no-cache-dir numexpr
+RUN pip install --no-cache-dir jupyterlab
+RUN pip install --no-cache-dir notebook
+RUN pip install --no-cache-dir scikit-learn
+RUN pip install --no-cache-dir scikit-optimize
+RUN pip install --no-cache-dir tensorflow
+RUN pip install --no-cache-dir xgboost
+RUN pip install --no-cache-dir scinum
+RUN pip install --no-cache-dir nvidia_smi
+RUN pip install --no-cache-dir py3nvml
+RUN pip install --no-cache-dir torch
 
 # install cmsml from master
 RUN git clone https://github.com/cms-ml/cmsml.git && \
diff --git a/docker/Dockerfile_39 b/docker/Dockerfile_39
index 98a152e..d78f27a 100644
--- a/docker/Dockerfile_39
+++ b/docker/Dockerfile_39
@@ -5,27 +5,26 @@ WORKDIR /root
 
 # minimal software stack
 RUN apt-get update; apt-get clean
-RUN apt-get install -y nano less htop git; apt-get clean
+RUN apt-get install -y nano less htop git libhdf5-serial-dev; apt-get clean
 
 # python software stack
 RUN pip install --no-cache-dir --upgrade pip setuptools
 RUN pip install --no-cache-dir --upgrade ipython
-RUN pip install --no-cache-dir \
-    numpy \
-    scipy \
-    matplotlib \
-    pandas \
-    numexpr \
-    jupyterlab \
-    notebook \
-    scikit-learn \
-    scikit-optimize \
-    tensorflow \
-    xgboost \
-    scinum \
-    nvidia_smi \
-    py3nvml \
-    torch
+RUN pip install --no-cache-dir numpy
+RUN pip install --no-cache-dir scipy
+RUN pip install --no-cache-dir matplotlib
+RUN pip install --no-cache-dir pandas
+RUN pip install --no-cache-dir numexpr
+RUN pip install --no-cache-dir jupyterlab
+RUN pip install --no-cache-dir notebook
+RUN pip install --no-cache-dir scikit-learn
+RUN pip install --no-cache-dir scikit-optimize
+RUN pip install --no-cache-dir tensorflow
+RUN pip install --no-cache-dir xgboost
+RUN pip install --no-cache-dir scinum
+RUN pip install --no-cache-dir nvidia_smi
+RUN pip install --no-cache-dir py3nvml
+RUN pip install --no-cache-dir torch
 
 # install cmsml from master
 RUN git clone https://github.com/cms-ml/cmsml.git && \
diff --git a/docker/Dockerfile_39_base b/docker/Dockerfile_39_base
index 035c135..e00979d 100644
--- a/docker/Dockerfile_39_base
+++ b/docker/Dockerfile_39_base
@@ -5,7 +5,7 @@ WORKDIR /root
 
 # minimal software stack
 RUN apt-get update; apt-get clean
-RUN apt-get install -y nano less htop git; apt-get clean
+RUN apt-get install -y nano less htop git libhdf5-serial-dev; apt-get clean
 
 # python software stack
 RUN pip install --no-cache-dir --upgrade pip setuptools
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 61c3cc4..0ab122d 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -1,5 +1,5 @@
-flake8~=5.0
-flake8-commas~=2.1
-flake8-quotes~=3.3
-pytest-cov>=3.0
+flake8~=7.0
+flake8-commas~=2.1.0
+flake8-quotes~=3.3.2
+pytest-cov~=5.0.0
 pytest-xdist~=3.4.0
diff --git a/tests/test_compile_tf_graph.py b/tests/test_compile_tf_graph.py
index f912921..73e28ab 100644
--- a/tests/test_compile_tf_graph.py
+++ b/tests/test_compile_tf_graph.py
@@ -38,11 +38,11 @@ def tf_version(self):
 
     def create_test_model(self, tf):
         # model with 2 input nodes and 1 output node, with non-static batchsize
-        x1 = tf.keras.Input(shape=(2,), name="first")
-        x2 = tf.keras.Input(shape=(3,), name="second")
-        x3 = tf.keras.Input(shape=(10,), name="third")
+        x1 = tf.keras.Input(shape=(2,), name="inputs")
+        x2 = tf.keras.Input(shape=(3,), name="inputs_1")
+        x3 = tf.keras.Input(shape=(10,), name="inputs_2")
 
-        x = tf.concat([x1, x2], axis=1)
+        x = tf.keras.layers.concatenate([x1, x2], axis=1)
         a1 = tf.keras.layers.Dense(10, activation="elu")(x)
         y = tf.keras.layers.Dense(5, activation="softmax")(a1)
 
@@ -93,20 +93,20 @@ def test_compile_tf_graph_static_preparation(self):
             model_static_inputs = loaded_static_model.signatures[key].structured_input_signature[1]
 
             expected_model_static_inputs = {
-                f"first_bs{batch_size}": tf.TensorSpec(
+                f"inputs_bs{batch_size}": tf.TensorSpec(
                     shape=(batch_size, 2),
                     dtype=tf.float32,
-                    name=f"first_bs{batch_size}",
+                    name=f"inputs_bs{batch_size}",
                 ),
-                f"second_bs{batch_size}": tf.TensorSpec(
+                f"inputs_1_bs{batch_size}": tf.TensorSpec(
                     shape=(batch_size, 3),
                     dtype=tf.float32,
-                    name=f"second_bs{batch_size}",
+                    name=f"inputs_1_bs{batch_size}",
                 ),
-                f"third_bs{batch_size}": tf.TensorSpec(
+                f"inputs_2_bs{batch_size}": tf.TensorSpec(
                     shape=(batch_size, 10),
                     dtype=tf.float32,
-                    name=f"third_bs{batch_size}",
+                    name=f"inputs_2_bs{batch_size}",
                 ),
             }
 
diff --git a/tests/test_tensorflow.py b/tests/test_tensorflow.py
index b2b3f00..c4d4fea 100644
--- a/tests/test_tensorflow.py
+++ b/tests/test_tensorflow.py
@@ -58,10 +58,15 @@ def b(self):
     def create_keras_model(self, tf):
         model = tf.keras.Sequential()
 
-        model.add(tf.keras.layers.InputLayer(input_shape=(10,), dtype=tf.float32, name="input"))
-        model.add(tf.keras.layers.BatchNormalization(axis=1, renorm=True))
+        # input_shape is deprected since TF 2.16
+        ge_2_16 = cmsml.tensorflow.tf_version_check(">=", (2, 16))
+        input_kwargs = {"dtype": tf.float32, "name": "input"}
+        input_kwargs["shape" if ge_2_16 else "input_shape"] = (10,)
+
+        model.add(tf.keras.layers.InputLayer(**input_kwargs))
+        model.add(tf.keras.layers.BatchNormalization(axis=1))
         model.add(tf.keras.layers.Dense(100, activation="tanh"))
-        model.add(tf.keras.layers.BatchNormalization(axis=1, renorm=True))
+        model.add(tf.keras.layers.BatchNormalization(axis=1))
         model.add(tf.keras.layers.Dense(3, activation="softmax", name="output"))
 
         return model
@@ -249,13 +254,14 @@ def test_save_keras_model_v1(self):
             cmsml.tensorflow.save_frozen_graph(path, model, variables_to_constants=True)
             self.assertTrue(os.path.exists(path))
 
-        with tmp_file(suffix=".pb") as path:
-            cmsml.tensorflow.save_frozen_graph(
-                path,
-                self.tf1.keras.backend.get_session(),
-                variables_to_constants=False,
-            )
-            self.assertTrue(os.path.exists(path))
+        if cmsml.tensorflow.tf_keras_version_check("<", 3):
+            with tmp_file(suffix=".pb") as path:
+                cmsml.tensorflow.save_frozen_graph(
+                    path,
+                    self.tf1.keras.backend.get_session(),
+                    variables_to_constants=False,
+                )
+                self.assertTrue(os.path.exists(path))
 
     def test_save_keras_model_v2(self):
         model = self.create_keras_model(self.tf)
@@ -339,9 +345,19 @@ def create_saved_model(self, **kwargs):
 
         model = self.create_keras_model(self.tf)
 
-        with tmp_dir(create=False) as keras_path, tmp_dir(create=False) as tf_path:
+        # keras 3 requires a file ending in .keras
+        if cmsml.tensorflow.tf_keras_version_check(">=", 3):
+            tmp_keras = lambda: tmp_file(create=False, suffix=".keras")
+        else:
+            tmp_keras = lambda: tmp_dir(create=False)
+
+        with tmp_keras() as keras_path, tmp_dir(create=False) as tf_path:
             self.tf.saved_model.save(model, tf_path)
-            model.save(keras_path, overwrite=True, include_optimizer=False)
+
+            keras_kwargs = {"overwrite": True}
+            if cmsml.tensorflow.tf_version_check("<=", (2, 11)):
+                keras_kwargs["include_optimizer"] = False
+            model.save(keras_path, **keras_kwargs)
 
             yield keras_path, tf_path
 
@@ -351,8 +367,13 @@ def test_load_model(self):
             keras_model = cmsml.tensorflow.load_model(keras_path)
             tf_model = cmsml.tensorflow.load_model(tf_path)
 
+            if cmsml.tensorflow.tf_version_check(">=", (2, 16)):
+                tf_func = lambda inp: tf_model.signatures["serving_default"](inp)["output_0"]
+            else:
+                tf_func = tf_model
+
             inp = self.tf.ones(shape=(2, 10))
-            keras_out, tf_out = keras_model(inp), tf_model(inp)
+            keras_out, tf_out = keras_model(inp), tf_func(inp)
 
             expected_shape = self.tf.TensorShape([2, 3])
 
@@ -363,8 +384,9 @@ def test_load_graph_def(self):
         with self.create_saved_model() as paths:
             keras_path, tf_path = paths
             default_serving_key = self.tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
-            tf_graph_def = cmsml.tensorflow.load_graph_def(tf_path, default_serving_key)
-            keras_graph_def = cmsml.tensorflow.load_graph_def(keras_path, default_serving_key)
 
+            tf_graph_def = cmsml.tensorflow.load_graph_def(tf_path, default_serving_key)
             self.assertTrue(isinstance(tf_graph_def, self.tf.compat.v1.GraphDef))
+
+            keras_graph_def = cmsml.tensorflow.load_graph_def(keras_path, default_serving_key)
             self.assertTrue(isinstance(keras_graph_def, self.tf.compat.v1.GraphDef))