diff --git a/.github/workflows/build-images.yml b/.github/workflows/build-images.yml
index d5cc93a..633ee96 100644
--- a/.github/workflows/build-images.yml
+++ b/.github/workflows/build-images.yml
@@ -37,7 +37,6 @@ jobs:
           - ubuntu2004_rocm_oneapi
           - ubuntu2004_oneapi
           - ubuntu2004_exatrkx
-          - ubuntu2004_exatrkx_training
           - ubuntu2204
           - ubuntu2204_clang
           - ubuntu2204_cpp20
diff --git a/ubuntu1804_cuda/Dockerfile b/ubuntu1804_cuda/Dockerfile
index 9c01ae2..54216a4 100644
--- a/ubuntu1804_cuda/Dockerfile
+++ b/ubuntu1804_cuda/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:10.2-devel-ubuntu18.04
+FROM nvidia/cuda:11.8.0-devel-ubuntu18.04
 
 LABEL description="Ubuntu 18.04 with Acts dependencies"
 LABEL maintainer="Paul Gessinger <paul.gessinger@cern.ch"
diff --git a/ubuntu2004_exatrkx/Dockerfile b/ubuntu2004_exatrkx/Dockerfile
index c4c5793..d518061 100644
--- a/ubuntu2004_exatrkx/Dockerfile
+++ b/ubuntu2004_exatrkx/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:11.5.1-cudnn8-devel-ubuntu20.04
+FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
 
 LABEL description="Ubuntu 20.04 with Acts dependencies for the Exa.TrkX Plugin"
 LABEL maintainer="Paul Gessinger <paul.gessinger@cern.ch"
@@ -83,8 +83,8 @@ RUN mkdir src \
   && rm -rf build src
 
 # libtorch (unzip cannot be used in a pipe...)
-ENV LIBTORCH_URL_GPU https://download.pytorch.org/libtorch/cu113/libtorch-cxx11-abi-shared-with-deps-1.10.2%2Bcu113.zip
-ENV LIBTORCH_URL_CPU https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.10.2%2Bcpu.zip
+ENV LIBTORCH_URL_GPU https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcu118.zip
+ENV LIBTORCH_URL_CPU https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.0.1%2Bcpu.zip
 
 RUN ${GET} --output libtorch.zip ${LIBTORCH_URL_GPU} \
   && unzip libtorch.zip \
diff --git a/ubuntu2004_exatrkx_training/Dockerfile b/ubuntu2004_exatrkx_training/Dockerfile
deleted file mode 100644
index 51311ef..0000000
--- a/ubuntu2004_exatrkx_training/Dockerfile
+++ /dev/null
@@ -1,32 +0,0 @@
-FROM nvidia/cuda:11.5.1-cudnn8-devel-ubuntu20.04
-
-LABEL description="Ubuntu 20.04 with prerequisites for running the Exa.TrkX training pipeline"
-LABEL maintainer="Benjamin Huth <benjamin.huth@ur.de>"
-# increase whenever any of the RUN commands change
-LABEL version="1"
-
-# DEBIAN_FRONTEND ensures non-blocking operation (tzdata is a problem)
-ENV DEBIAN_FRONTEND noninteractive
-
-# install dependencies from the package manager
-RUN apt-get update -y \
-  && apt-get install -y \
-    git \
-    python3 \
-    python3-dev \
-    python3-pip \
-  && apt-get clean -y
-  
-# Software for the training
-RUN pip3 install torch==1.11.0+cu115 --extra-index-url https://download.pytorch.org/whl/cu115
-RUN pip3 install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.11.0+cu115.html
-RUN pip3 install traintrack==0.1.5
-RUN pip3 install cupy==11.1.0
-RUN pip3 install faiss-gpu==1.7.2
-
-ENV TORCH_CUDA_ARCH_LIST="7.5;8.6"
-RUN git clone --recursive https://github.com/lxxue/FRNN src \
-  && (cd src && git checkout 3e370d8d9073d4e130363faf87d2370598b5fbf2) \
-  && (cd src/external/prefix_sum && pip install .) \
-  && (cd src && pip install .) \
-  && rm -rf src