diff --git a/.github/container/test-pax.sh b/.github/container/test-pax.sh
index 5562b729a..91cb926b4 100755
--- a/.github/container/test-pax.sh
+++ b/.github/container/test-pax.sh
@@ -388,6 +388,9 @@ else
 fi
 
 if [[ ${EVALUATE} -ne 0 ]]; then
+
+  trap "rm -rf ${OUTPUT}/checkpoints" ERR INT HUP TERM EXIT
+
   ## train for 0 steps to generate an initial checkpoint
   python -m paxml.main \
     --fdl_config=${CONFIG} \
@@ -408,7 +411,6 @@ if [[ ${EVALUATE} -ne 0 ]]; then
     $ADDITIONAL_ARGS \
     $([[ $MULTIPROCESS != 0 ]] && echo --multiprocess_gpu)
 
-  rm -rf ${OUTPUT}/checkpoints
 else
   python -m paxml.main \
     --fdl_config=${CONFIG} \