diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index c5e4934..9fd48aa 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -18,11 +18,14 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v2
         with:
-          python-version: '3.9'
+          python-version: '3.11'
 
       - name: Install dependencies
         run: pip install -r requirements.txt
 
+      - name: Install pytest
+        run: pip install pytest
+
       - name: Set PYTHONPATH to subdirectory sdcat
         run: echo "PYTHONPATH=." >> $GITHUB_ENV
 
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 0af2e3e..ccda88c 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -26,7 +26,7 @@ jobs:
           docker login -u ${{ secrets.DOCKERHUB_USERNAME }} -p ${{ secrets.DOCKERHUB_TOKEN }} 
           docker buildx create --name mybuilder --platform linux/amd64,linux/arm64 --use
           cd docker && \
-          docker buildx build --push --platform linux/amd64 -t mbari/sdcat:$RELEASE_VERSION-cuda12.4.1  --label IMAGE_URI=mbari/sdcat:$RELEASE_VERSION-cuda12.4.1 -f Dockerfile.cuda12 .
+          docker buildx build --push --platform linux/amd64 -t mbari/sdcat:$RELEASE_VERSION-cuda12  --label IMAGE_URI=mbari/sdcat:$RELEASE_VERSION-cuda12 -f Dockerfile.cuda12 .
           docker buildx build --push --platform linux/amd64,linux/arm64 -t mbari/sdcat:$RELEASE_VERSION  --label IMAGE_URI=mbari/sdcat:$RELEASE_VERSION -f Dockerfile .
   push_readme_to_dockerhub:
     runs-on: ubuntu-latest
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 6e66bd3..e0b5cfe 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -31,8 +31,10 @@ RUN if [ "$GIT_VERSION" != "latest" ]; then \
     git clone --depth 1 https://github.com/mbari-org/sdcat.git;\
     fi
 
-WORKDIR /home/docker_user/sdcat
-RUN pip install -r requirements.txt
+RUN sed -i '/torch==2.3.1/d' /home/docker_user/sdcat/requirements.txt && \
+    sed -i '/torchvision==0.18.1/d' /home/docker_user/sdcat/requirements.txt \
+    && pip install -r /home/docker_user/sdcat/requirements.txt \
+    && pip install torchvision==0.18.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
 
 WORKDIR /tmp
 RUN git clone https://github.com/DmitryUlyanov/Multicore-TSNE.git && \
diff --git a/docker/Dockerfile.cuda12 b/docker/Dockerfile.cuda12
index 8297b74..d65c8e3 100644
--- a/docker/Dockerfile.cuda12
+++ b/docker/Dockerfile.cuda12
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
+FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime
 
 LABEL vendor="MBARI"
 LABEL maintainer="dcline@mbari.org"
@@ -12,22 +12,17 @@ RUN apt-get update && apt-get install -y \
     git \
     wget \
     curl \
-    && rm -rf /var/lib/apt/lists/*
-
-RUN apt-get update && apt-get install -y \
-    python3 \
-    python3-pip \
-    python3.10-venv \
+    libgl1 \
+    libglib2.0-0 \
     && rm -rf /var/lib/apt/lists/*
 
 ARG GIT_VERSION=latest
 ARG IMAGE_URI=mbari/sdcat:${GIT_VERSION}
-ARG DOCKER_GID=1001
-ARG DOCKER_UID=12065
+ARG DOCKER_GID=20
+ARG DOCKER_UID=12078
 
 # Create a user with a specific UID and GID.
-RUN addgroup --gid ${DOCKER_GID} docker && \
-    adduser  --uid ${DOCKER_UID} --gid ${DOCKER_GID} --disabled-password --shell /bin/bash --gecos "Docker" docker_user
+RUN adduser  --uid ${DOCKER_UID} --gid ${DOCKER_GID} --disabled-password --shell /bin/bash --gecos "Docker" docker_user
 
 # Install in the virtual environment as the docker user.
 USER docker_user
@@ -44,14 +39,29 @@ RUN if [ "$GIT_VERSION" != "latest" ]; then \
     git clone --depth 1 https://github.com/mbari-org/sdcat.git;\
     fi
 
-WORKDIR /home/docker_user/sdcat
-RUN pip install -r requirements.txt
+RUN sed -i '/torch==2.3.1/d' /home/docker_user/sdcat/requirements.txt && \
+    sed -i '/torchvision==0.18.1/d' /home/docker_user/sdcat/requirements.txt \
+    && pip install -r /home/docker_user/sdcat/requirements.txt
 
 WORKDIR /tmp
 RUN git clone https://github.com/DmitryUlyanov/Multicore-TSNE.git && \
     cd Multicore-TSNE && \
     pip install .  && \
     cd .. && \
-    rm -rf Multicore-TSNE \
+    rm -rf Multicore-TSNE
+
+# Remove the temporary directory and other build tools to reduce the image size
+USER root
+RUN apt-get remove -y \
+    build-essential \
+    cmake \
+    git \
+    wget \
+    curl \
+    && rm -rf /var/cache/apk/* \
+    && rm -rf /home/docker_user/sdcat/.git \
+    && rm -rf /home/docker_user/sdcat/sdcat/tests
 
-ENTRYPOINT ["python", "sdcat"]
\ No newline at end of file
+USER docker_user
+WORKDIR /home/docker_user/sdcat
+ENTRYPOINT ["python", "-m", "sdcat"
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 5694c4b..7d983a0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,26 +6,22 @@ mmdet==2.21.0
 mmengine==0.7.2
 click
 hdbscan==0.8.33
-torch==2.2.1
+torch==2.3.1
 tator
 piexif
 yolov5==7.0.13
-torchvision==0.17.1
+torchvision==0.18.1
 transformers
 timm
 pandas>=1.2.4
 ultralytics
-python-dotenv
 pascal-voc-writer
 scikit-learn==1.4.0rc1
 scikit-image
-plotly
-pytest
 huggingface-hub==0.19.4
 umap-learn==0.5.5
 hdbscan>=0.8.27
 matplotlib==3.7.0
 numba==0.57
 ephem
-jupyter
 tqdm