diff --git a/.github/workflows/build-container.yaml b/.github/workflows/build-container.yaml
index 86e17fb..41fec71 100644
--- a/.github/workflows/build-container.yaml
+++ b/.github/workflows/build-container.yaml
@@ -14,6 +14,7 @@ on:
       - pyproject.toml
       - requirements-gpu.txt
       - .github/workflows/build-container.yaml
+      - .github/workflows/compute-tag.yaml
 
 jobs:
   compute_tag:
diff --git a/.github/workflows/compute-tag.yaml b/.github/workflows/compute-tag.yaml
index 37fcbfc..805ce72 100644
--- a/.github/workflows/compute-tag.yaml
+++ b/.github/workflows/compute-tag.yaml
@@ -2,11 +2,18 @@ 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
@@ -20,7 +27,17 @@ jobs:
             echo "tag=latest" >> $GITHUB_OUTPUT
           elif [[ "${{ github.event_name }}" == "pull_request" ]]; then
             sanitized_branch_name="${branch_name//\//-}"
-            echo "tag=${sanitized_branch_name}" >> $GITHUB_OUTPUT
+            # 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 beisner/taxpose:${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//\//-}"
diff --git a/.github/workflows/merge-request.yaml b/.github/workflows/merge-request.yaml
index 373ac65..d33dcf3 100644
--- a/.github/workflows/merge-request.yaml
+++ b/.github/workflows/merge-request.yaml
@@ -10,6 +10,8 @@ on:
 jobs:
   compute_tag:
     uses: ./.github/workflows/compute-tag.yaml
+    with:
+      latest_on_noexist: 'true'
 
   test:
     uses: ./.github/workflows/run-tests.yaml
diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml
index dacbbd7..1599bc8 100644
--- a/.github/workflows/run-tests.yaml
+++ b/.github/workflows/run-tests.yaml
@@ -25,10 +25,16 @@ jobs:
         working-directory: /opt/baeisner/code
 
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           submodules: 'true'
 
+      # 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 taxpose/ tests/ --check
 
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 3c942fb..4df8057 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -10,6 +10,22 @@ pre-commit install
 ## Run CI locally
 To run the CI locally:
 
+### [optional] Start the cache.
+
+Download
+https://github.com/sp-ricard-valverde/github-act-cache-server
+
+Run
+```
+ACT_CACHE_AUTH_KEY=foo docker compose up --build
+```
+
+Also modify your .actrc.
+
+
+### Run
+
+
 Setup (make sure docker is installed):
 ```
 brew install act
diff --git a/taxpose/training/pm_baselines/test_bc.py b/taxpose/training/pm_baselines/test_bc.py
index 8812599..81e715c 100644
--- a/taxpose/training/pm_baselines/test_bc.py
+++ b/taxpose/training/pm_baselines/test_bc.py
@@ -252,7 +252,7 @@ def load_bc_model(method: str, exp_name: str) -> DaggerNet:
     net = DaggerNet.load_from_checkpoint(
         f"{d}/{ckpt}",
     )
-    return net.cuda()
+    return net.cuda()  # type: ignore
 
 
 def predict_next_step_bc(