diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 52138998c45..d3325fd503c 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -55,94 +55,96 @@ clang_tidy:
     # Generate an html report
     - codereport clang-tidy/clang-tidy.json clang-tidy/html
 
-# build_exatrkx_cpu:
-#   stage: build
-#   image: ghcr.io/acts-project/ubuntu2204_exatrkx:52
-#   tags:
-#     - large
+build_exatrkx_cpu:
+  stage: build
+  image: ghcr.io/acts-project/ubuntu2204_exatrkx:52
+  tags:
+    - large
 
-#   cache:
-#     key: ccache-exatrkx-cpu-$CI_COMMIT_REF_SLUG
-#     paths:
-#       - ${CI_PROJECT_DIR}/ccache
+  cache:
+    key: ccache-exatrkx-cpu-$CI_COMMIT_REF_SLUG
+    paths:
+      - ${CI_PROJECT_DIR}/ccache
 
-#   script:
-#     - export PATH=/usr/local/sbin:/usr/sbin:/sbin:$PATH
-#     - export PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
-#     - echo $PATH
-#     - git clone $CLONE_URL src
-#     - cd src
-#     - git checkout $HEAD_SHA
-#     - cd ..
-#     - mkdir build
-#     # Here we only do a minimal build without examples to save ressources
-#     - >
-#       cmake -B build -S src
-#       -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
-#       -GNinja
-#       -DCMAKE_BUILD_TYPE=Release
-#       -DCMAKE_CXX_FLAGS=-w
-#       -DACTS_BUILD_PLUGIN_EXATRKX=ON
-#       -DACTS_BUILD_EXAMPLES_EXATRKX=OFF
-#       -DACTS_BUILD_UNITTESTS=OFF
-#       -DACTS_EXATRKX_ENABLE_TORCH=ON
-#       -DACTS_EXATRKX_ENABLE_ONNX=OFF
-#       -DACTS_EXATRKX_ENABLE_CUDA=OFF
-#       -DACTS_BUILD_EXAMPLES_PYTHON_BINDINGS=OFF
-#       -DACTS_ENABLE_LOG_FAILURE_THRESHOLD=ON
-
-#     - ccache -z
-#     - cmake --build build -- -j2
-#     - ccache -s
-
-# build_exatrkx:
-#   stage: build
-#   image: ghcr.io/acts-project/ubuntu2204_exatrkx:52
-#   tags:
-#     - large
-#
-#   cache:
-#     key: ccache-exatrkx-$CI_COMMIT_REF_SLUG
-#     paths:
-#       - ${CI_PROJECT_DIR}/ccache
-#
-#   artifacts:
-#     paths:
-#       - build/
-#     exclude:
-#       - build/**/*.o
-#       - build/bin/ActsIntegrationTest*
-#     expire_in: 6 hours
-#
-#   script:
-#     - export PATH=/usr/local/sbin:/usr/sbin:/sbin:$PATH
-#     - export PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
-#     - echo $PATH
-#     - git clone $CLONE_URL src
-#     - cd src
-#     - git checkout $HEAD_SHA
-#     - cd ..
-#     - mkdir build
-#     - >
-#       cmake -B build -S src
-#       -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
-#       -GNinja
-#       -DCMAKE_BUILD_TYPE=Release
-#       -DCMAKE_CXX_FLAGS=-w
-#       -DCMAKE_CUDA_FLAGS=-w
-#       -DCMAKE_CUDA_ARCHITECTURES="75;86"
-#       -DACTS_BUILD_PLUGIN_EXATRKX=ON
-#       -DACTS_BUILD_EXAMPLES_EXATRKX=ON
-#       -DACTS_BUILD_UNITTESTS=ON
-#       -DACTS_EXATRKX_ENABLE_TORCH=ON
-#       -DACTS_EXATRKX_ENABLE_ONNX=ON
-#       -DACTS_EXATRKX_ENABLE_CUDA=ON
-#       -DACTS_BUILD_EXAMPLES_PYTHON_BINDINGS=ON
-#       -DACTS_ENABLE_LOG_FAILURE_THRESHOLD=ON
-#
-#     - ccache -z
-#     - cmake --build build -- -j2
-#     - ccache -s
+  script:
+    - export PATH=/usr/local/sbin:/usr/sbin:/sbin:$PATH
+    - export PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
+    - export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
+    - echo $PATH
+    - git clone $CLONE_URL src
+    - cd src
+    - git checkout $HEAD_SHA
+    - cd ..
+    - mkdir build
+    # Here we only do a minimal build without examples to save ressources
+    - >
+      cmake -B build -S src
+      -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
+      -GNinja
+      -DCMAKE_BUILD_TYPE=Release
+      -DCMAKE_CXX_FLAGS=-w
+      -DACTS_BUILD_PLUGIN_EXATRKX=ON
+      -DACTS_BUILD_EXAMPLES_EXATRKX=OFF
+      -DACTS_BUILD_UNITTESTS=OFF
+      -DACTS_EXATRKX_ENABLE_TORCH=ON
+      -DACTS_EXATRKX_ENABLE_ONNX=OFF
+      -DACTS_EXATRKX_ENABLE_CUDA=OFF
+      -DACTS_BUILD_EXAMPLES_PYTHON_BINDINGS=OFF
+      -DACTS_ENABLE_LOG_FAILURE_THRESHOLD=ON
+
+    - ccache -z
+    - cmake --build build -- -j2
+    - ccache -s
+
+build_exatrkx:
+  stage: build
+  image: ghcr.io/acts-project/ubuntu2204_exatrkx:52
+  tags:
+    - large
+
+  cache:
+    key: ccache-exatrkx-$CI_COMMIT_REF_SLUG
+    paths:
+      - ${CI_PROJECT_DIR}/ccache
+
+  artifacts:
+    paths:
+      - build/
+    exclude:
+      - build/**/*.o
+      - build/bin/ActsIntegrationTest*
+    expire_in: 6 hours
+
+  script:
+    - export PATH=/usr/local/sbin:/usr/sbin:/sbin:$PATH
+    - export PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
+    - export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
+    - echo $PATH
+    - git clone $CLONE_URL src
+    - cd src
+    - git checkout $HEAD_SHA
+    - cd ..
+    - mkdir build
+    - >
+      cmake -B build -S src
+      -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
+      -GNinja
+      -DCMAKE_BUILD_TYPE=Release
+      -DCMAKE_CXX_FLAGS=-w
+      -DCMAKE_CUDA_FLAGS=-w
+      -DCMAKE_CUDA_ARCHITECTURES="75;86"
+      -DACTS_BUILD_PLUGIN_EXATRKX=ON
+      -DACTS_BUILD_EXAMPLES_EXATRKX=ON
+      -DACTS_BUILD_UNITTESTS=ON
+      -DACTS_EXATRKX_ENABLE_TORCH=ON
+      -DACTS_EXATRKX_ENABLE_ONNX=OFF
+      -DACTS_EXATRKX_ENABLE_CUDA=ON
+      -DACTS_BUILD_EXAMPLES_PYTHON_BINDINGS=ON
+      -DACTS_ENABLE_LOG_FAILURE_THRESHOLD=ON
+
+    - ccache -z
+    - cmake --build build -- -j2
+    - ccache -s
 
 # test_exatrkx_unittests:
 #   stage: test