diff --git a/.github/workflows/tpu_unit_tests.yaml b/.github/workflows/tpu_unit_tests.yaml
index a2d4c8958..84b7c633c 100644
--- a/.github/workflows/tpu_unit_tests.yaml
+++ b/.github/workflows/tpu_unit_tests.yaml
@@ -32,4 +32,9 @@ jobs:
           export TPU_NAME=ci-run-${{ github.run_id }}
           eval "$(ssh-agent -s)"
           infra/babysit-tpu-vm.sh $TPU_NAME -z ${{ matrix.tpu-zone }} -t v4-8 --preemptible -s infra/helpers/setup-tpu-vm-tests.sh -b ${{ github.sha }} --retries 3 -- \
-            bash levanter/infra/run.sh pytest levanter/tests
+            PYTHONPATH=$PYTHONPATH:levanter/tests bash levanter/infra/run.sh pytest levanter/tests
+
+      - name: Cleanup
+        if: ${{ always() }}
+        run: |
+          gcloud compute tpus tpu-vm delete -y $TPU_NAME --zone ${{ matrix.tpu-zone }} --quiet