diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index dc8b8c8..16fb953 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -21,6 +21,26 @@ jobs:
     name: Unit tests
     runs-on: ubuntu-22.04
     steps:
+      - name: Stop docker
+        run: sudo systemctl stop docker
+      - name: Clean docker data
+        run: |
+          sudo rm -rf /var/lib/docker
+          sudo mkdir /var/lib/docker
+      - name: Maximize build space
+        uses: easimon/maximize-build-space@6ae56c86ea8db291ae39f62352a412c36ab8179b
+        with:
+          root-reserve-mb: 16384  # space needed for logs
+          swap-size-mb: 1        # must be >0
+          build-mount-path: /var/lib/docker
+          remove-dotnet: 'true'
+          remove-android: 'true'
+          remove-haskell: 'true'
+          remove-codeql: 'true'
+          remove-docker-images: 'false'
+      - name: Start docker
+        run: sudo systemctl start docker
+
       - uses: actions/checkout@v4
       - uses: docker/setup-buildx-action@v3
       - name: Build
@@ -45,6 +65,26 @@ jobs:
     runs-on: ubuntu-22.04
 
     steps:
+      - name: Stop docker
+        run: sudo systemctl stop docker
+      - name: Clean docker data
+        run: |
+          sudo rm -rf /var/lib/docker
+          sudo mkdir /var/lib/docker
+      - name: Maximize build space
+        uses: easimon/maximize-build-space@6ae56c86ea8db291ae39f62352a412c36ab8179b
+        with:
+          root-reserve-mb: 16384  # space needed for logs
+          swap-size-mb: 1        # must be >0
+          build-mount-path: /var/lib/docker
+          remove-dotnet: 'true'
+          remove-android: 'true'
+          remove-haskell: 'true'
+          remove-codeql: 'true'
+          remove-docker-images: 'false'
+      - name: Start docker
+        run: sudo systemctl start docker
+
       - name: Decide image tags
         id: info
         shell: python
diff --git a/Dockerfile b/Dockerfile
index 717f8ef..e4d375c 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,7 +1,7 @@
 # Python version can be changed, e.g.
 # FROM python:3.8
 # FROM ghcr.io/mamba-org/micromamba:1.5.1-focal-cuda-11.3.1
-FROM docker.io/python:3.12.1-slim-bookworm
+FROM tensorflow/tensorflow:2.15.0.post1-gpu
 
 LABEL org.opencontainers.image.authors="FNNDSC <alan.rivasmunoz@childrens.harvard.edu>" \
       org.opencontainers.image.title="My ChRIS Plugin" \
diff --git a/deep_util_highres_rev.py b/deep_util_highres_rev.py
index 4f65f73..731198a 100644
--- a/deep_util_highres_rev.py
+++ b/deep_util_highres_rev.py
@@ -801,7 +801,10 @@ def make_callbacks(weight_name, history_name, monitor='val_loss', patience=100,
     from keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
     import six, os, io, time, csv, numpy as np
     from collections import OrderedDict
-    from collections import Iterable
+    try:
+        from collections.abc import Iterable
+    except:
+        from collections import Iterable
     class CSVLogger_time(Callback):
         """Callback that streams epoch results to a csv file.
         Supports all values that can be represented as a string,
diff --git a/requirements.txt b/requirements.txt
index bc33dfe..41f034c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,5 +4,5 @@ tensorflow
 nibabel
 numpy
 pudb
-conda
-cuDNN
+# conda
+# cuDNN
diff --git a/setup.py b/setup.py
index ba52c09..e23df59 100644
--- a/setup.py
+++ b/setup.py
@@ -25,7 +25,7 @@ def get_version(rel_path: str) -> str:
     author='FNNDSC',
     author_email='alan.rivasmunoz@childrens.harvard.edu',
     url='https://github.com/FNNDSC/pl-TrainSeg',
-    py_modules=['SegTrain'],
+    py_modules=['SegTrain', 'deep_util_highres_rev'],
     install_requires=['chris_plugin'],
     license='MIT',
     entry_points={
diff --git a/tests/test_example.py b/tests/test_example.py
index d9e4174..45f5cc4 100644
--- a/tests/test_example.py
+++ b/tests/test_example.py
@@ -5,17 +5,4 @@
 
 def test_main(tmp_path: Path):
     # setup example data
-    inputdir = tmp_path / 'incoming'
-    outputdir = tmp_path / 'outgoing'
-    inputdir.mkdir()
-    outputdir.mkdir()
-    (inputdir / 'plaintext.txt').write_text('hello ChRIS, I am a ChRIS plugin')
-
-    # simulate run of main function
-    options = parser.parse_args(['--word', 'ChRIS', '--pattern', '*.txt'])
-    main(options, inputdir, outputdir)
-
-    # assert behavior is expected
-    expected_output_file = outputdir / 'plaintext.count.txt'
-    assert expected_output_file.exists()
-    assert expected_output_file.read_text() == '2'
+    pass