diff --git a/.github/workflows/build-container.yaml b/.github/workflows/build-container.yaml
new file mode 100644
index 0000000..44d878a
--- /dev/null
+++ b/.github/workflows/build-container.yaml
@@ -0,0 +1,43 @@
+# A github action that builds a container image for the project.
+
+name: Build Container
+
+on:
+  push:
+
+    # Allows you to run this workflow manually from the Actions tab
+    workflow_dispatch:
+
+    paths:
+      # This is the entire list of files that will trigger the workflow.
+      - Dockerfile
+      - pyproject.toml
+      - requirements-gpu.txt
+      - .github/workflows/build-container.yaml
+      - .github/workflows/compute-tag.yaml
+
+jobs:
+  compute_tag:
+    uses: ./.github/workflows/compute-tag.yaml
+
+  docker:
+    runs-on: ubuntu-latest
+    needs: compute_tag
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          push: true
+          # This is the name of the image that will be pushed to Docker Hub. If the branch is main, the image will be tagged as latest. Else, it will be tagged as the branch name.
+          tags: ${{ secrets.DOCKERHUB_USERNAME }}/python_ml_project_template:${{ needs.compute_tag.outputs.image_tag }}
+          cache-from: type=registry,ref=${{ secrets.DOCKERHUB_USERNAME }}/python_ml_project_template:${{ needs.compute_tag.outputs.image_tag }}
+          cache-to: type=inline
diff --git a/.github/workflows/build-site.yaml b/.github/workflows/build-site.yaml
index eabaf89..3e47f36 100644
--- a/.github/workflows/build-site.yaml
+++ b/.github/workflows/build-site.yaml
@@ -7,31 +7,36 @@ jobs:
   build-docs:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
+        with:
+          submodules: 'true'
 
       ##############################################
       # Skip caching if using a local runner.
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         if: ${{ !env.ACT }}
         with:
           python-version: '3.10'
           cache: 'pip'
           cache-dependency-path: "pyproject.toml"
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         if: ${{ env.ACT }}
         with:
           python-version: '3.10'
       ##############################################
 
-      - name: Install Dependencies
-        run: pip install -e ".[build_docs]"
+      - name: Install specific pip.
+        run: pip install pip==23.0.0
+
+      - name: Install doc requirements.
+        run: pip install mkdocs-material mkdocstrings[python]
 
       - name: Build mkdocs site
         working-directory: docs
         run: mkdocs build
 
       - name: Upload the built site.
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         if: ${{ !env.ACT }}
         with:
           name: site
diff --git a/.github/workflows/compute-tag.yaml b/.github/workflows/compute-tag.yaml
new file mode 100644
index 0000000..fccf6f2
--- /dev/null
+++ b/.github/workflows/compute-tag.yaml
@@ -0,0 +1,45 @@
+name: Compute the docker tag for this branch
+
+on:
+  workflow_call:
+    inputs:
+      # description: 'If true, the tag will be latest if the docker image tag does not exist'
+      latest_on_noexist:
+        required: false
+        type: string
+        default: 'false'
+    outputs:
+      image_tag:
+        description: 'The tag to use for the docker image'
+        value: ${{ jobs.compute_tag.outputs.image_tag }}
+
+
+jobs:
+  compute_tag:
+    runs-on: ubuntu-latest
+    outputs:
+      image_tag: ${{ steps.set_tag.outputs.tag }}
+    steps:
+      - id: set_tag
+        run: |
+          branch_name="${{ github.head_ref }}"
+          if [[ "${{ github.ref }}" == "refs/heads/main" ]]; then
+            echo "tag=latest" >> $GITHUB_OUTPUT
+          elif [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            sanitized_branch_name="${branch_name//\//-}"
+            # If latest_on_noexist is true, set the tag to latest if the tag does not exist.
+            if [[ "${{ inputs.latest_on_noexist }}" == "true" ]]; then
+              # Check if the tag exists using docker manifest.
+              if ! docker manifest inspect ${{ secrets.DOCKERHUB_USERNAME }}/python_ml_project_template:${sanitized_branch_name} > /dev/null 2>&1; then
+                echo "tag=latest" >> $GITHUB_OUTPUT
+              else
+                echo "tag=${sanitized_branch_name}" >> $GITHUB_OUTPUT
+              fi
+            else
+              echo "tag=${sanitized_branch_name}" >> $GITHUB_OUTPUT
+            fi
+          else
+            sanitized_branch_name="${GITHUB_REF#refs/heads/}"
+            sanitized_branch_name="${sanitized_branch_name//\//-}"
+            echo "tag=${sanitized_branch_name}" >> $GITHUB_OUTPUT
+          fi
diff --git a/.github/workflows/deploy-site.yaml b/.github/workflows/deploy-site.yaml
index 98dd7a1..fb275cb 100644
--- a/.github/workflows/deploy-site.yaml
+++ b/.github/workflows/deploy-site.yaml
@@ -18,18 +18,18 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Download Site Artifact
-        uses: actions/download-artifact@v3
+        uses: actions/download-artifact@v4
         with:
           name: site
           path: docs/site/
 
       - name: Setup Pages
         if: ${{ !env.ACT }}
-        uses: actions/configure-pages@v1
+        uses: actions/configure-pages@v5
 
       - name: Upload Artifact to Pages
         if: ${{ !env.ACT }}
-        uses: actions/upload-pages-artifact@v1
+        uses: actions/upload-pages-artifact@v3
         with:
           path: docs/site/
 
diff --git a/.github/workflows/merge-request.yaml b/.github/workflows/merge-request.yaml
index 8e227ad..d33dcf3 100644
--- a/.github/workflows/merge-request.yaml
+++ b/.github/workflows/merge-request.yaml
@@ -8,9 +8,18 @@ on:
   workflow_dispatch:
 
 jobs:
+  compute_tag:
+    uses: ./.github/workflows/compute-tag.yaml
+    with:
+      latest_on_noexist: 'true'
+
   test:
     uses: ./.github/workflows/run-tests.yaml
+    needs: compute_tag
     with:
       install_string: .[develop]
+      # Get the image tag from the compute_tag job.
+      image_tag: ${{ needs.compute_tag.outputs.image_tag }}
+
   build_site:
     uses: ./.github/workflows/build-site.yaml
diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml
index 2b14daa..857c021 100644
--- a/.github/workflows/push.yaml
+++ b/.github/workflows/push.yaml
@@ -8,10 +8,18 @@ on:
   workflow_dispatch:
 
 jobs:
+  compute_tag:
+    uses: ./.github/workflows/compute-tag.yaml
+    with:
+      latest_on_noexist: 'true'
   test:
     uses: ./.github/workflows/run-tests.yaml
+    needs: compute_tag
     with:
       install_string: .[develop]
+      # Get the image tag from the compute_tag job.
+      image_tag: ${{ needs.compute_tag.outputs.image_tag }}
+
   build_site:
     uses: ./.github/workflows/build-site.yaml
   deploy_site:
diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml
index 5426c1d..797e0f7 100644
--- a/.github/workflows/run-tests.yaml
+++ b/.github/workflows/run-tests.yaml
@@ -6,30 +6,33 @@ on:
       install_string:
         required: True
         type: string
+      image_tag:
+        required: True
+        type: string
+        default: "latest"
 
 jobs:
   test:
     runs-on: ubuntu-latest
 
-    steps:
-      - uses: actions/checkout@v3
+    container:
+      # Image tag is "latest" if the branch is main, else it is the branch name.
+      image: beisner/python_ml_project_template:${{ inputs.image_tag }}
 
-      ##############################################
-      # Skip caching if using a local runner.
-      - uses: actions/setup-python@v4
-        if: ${{ !env.ACT }}
-        with:
-          python-version: '3.10'
-          cache: 'pip'
-          cache-dependency-path: "pyproject.toml"
-      - uses: actions/setup-python@v4
-        if: ${{ env.ACT }}
+    defaults:
+      run:
+        working-directory: /opt/baeisner/code
+
+    steps:
+      - uses: actions/checkout@v4
         with:
-          python-version: '3.10'
-      ##############################################
+          submodules: 'true'
 
-      - name: Install package
-        run: pip install "${{ inputs.install_string }}"
+      # Link the code from the default checkout directory to the correct directory.
+      # Use the github workspace variable to get the correct directory.
+      # Can't use the checkout action to checkout to a different directory, so we have to simlink.
+      - name: Move code to correct directory
+        run: rm -rf /opt/baeisner/code && ln -s $GITHUB_WORKSPACE /opt/baeisner/code
 
       - name: Code Quality
         run: python -m black src/ tests/ --check
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 0c4b36b..9376e6b 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -2,7 +2,7 @@
     "editor.formatOnSave": true,
     "python.formatting.provider": "none",
     "editor.codeActionsOnSave": {
-        "source.organizeImports": true
+        "source.organizeImports": "explicit"
     },
     "[python]": {
         "editor.defaultFormatter": "ms-python.black-formatter"
diff --git a/Dockerfile b/Dockerfile
index b6be688..56b5814 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -37,7 +37,7 @@ WORKDIR $CODING_ROOT/code
 COPY ./src $CODING_ROOT/code/src
 COPY ./setup.py $CODING_ROOT/code/setup.py
 COPY ./pyproject.toml $CODING_ROOT/code/pyproject.toml
-RUN pip install -e .
+RUN pip install -e .[develop]
 
 # Changes to the configs and scripts will not require a rebuild
 COPY ./configs $CODING_ROOT/code/configs
diff --git a/README.md b/README.md
index 44dab1d..9c24e63 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +78,16 @@ To push this:
 docker push <my_dockerhub_username>/python-ml-project-template:latest
 ```
 
+## Using the CI.
+
+Set up pushing to docker:
+
+Put the following secrets in the Github repository:
+* `DOCKERHUB_USERNAME`: Your Dockerhub username
+* `DOCKERHUB_TOKEN`: Your Dockerhub token
+
+You'll also need to Ctrl-F replace instances of beisner and baeisner with appropriate usernames.
+
 ## Running on Clusters
 
 * [Autobot](autobot.md)
diff --git a/pyproject.toml b/pyproject.toml
index 523127f..e0dcf4b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,27 +4,21 @@ version = "0.1.0"
 description = "A Python Package Template"
 readme = "README.md"
 requires-python = ">=3.6"
-license = {file = "LICENSE.txt"}
-authors = [
-  {email = "baeisner@andrew.cmu.edu", name = "Ben Eisner"}
-]
+license = { file = "LICENSE.txt" }
+authors = [{ email = "baeisner@andrew.cmu.edu", name = "Ben Eisner" }]
 dependencies = [
   "hydra-core == 1.3.2",
   "lightning == 2.0.3",
   "omegaconf == 2.3.0",
   "pandas",
-  "torch == 2.0.1", # CUDA 11.8
+  "torch == 2.0.1",        # CUDA 11.8
   "torchmetrics",
   "torchvision == 0.15.2", # CUDA 11.8
   "wandb == 0.15.4",
 ]
 
 [build-system]
-requires = [
-  "setuptools >= 62.3.2",
-  "setuptools-scm",
-  "wheel",
-]
+requires = ["setuptools >= 62.3.2", "setuptools-scm", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project.optional-dependencies]
@@ -38,13 +32,8 @@ develop = [
   "pytest == 7.3.2",
   "pre-commit == 3.3.3",
 ]
-notebooks = [
-  "jupyter",
-]
-build_docs = [
-  "mkdocs-material",
-  "mkdocstrings[python]",
-]
+notebooks = ["jupyter"]
+build_docs = ["mkdocs-material", "mkdocstrings[python]"]
 
 # This is required to allow us to have notebooks/ at the top level.
 [tool.setuptools.packages.find]
@@ -58,7 +47,7 @@ profile = "black"
 known_third_party = "wandb"
 
 [tool.mypy]
-python_version = 3.8
+python_version = "3.10"
 warn_return_any = true
 warn_unused_configs = true
 mypy_path = "src"
@@ -66,11 +55,12 @@ namespace_packages = true
 explicit_package_bases = true
 
 [[tool.mypy.overrides]]
-module = [
-  "torchvision.*",
-]
+module = ["torchvision.*"]
 ignore_missing_imports = true
 
+[tool.pytest.ini_options]
+testpaths = "tests"
+
 [tool.pylint]
 known-third-party = "wandb"