diff --git a/.github/workflows/docker-publish-f5-tts.yml b/.github/workflows/docker-publish-f5-tts.yml
index 69ecbdb..6204812 100644
--- a/.github/workflows/docker-publish-f5-tts.yml
+++ b/.github/workflows/docker-publish-f5-tts.yml
@@ -32,12 +32,13 @@ jobs:
           password: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Build image
-        run: cd tts/F5-TTS && docker build . --file Dockerfile --tag $IMAGE_NAME --label "runnumber=${GITHUB_RUN_ID}"
-
-      - name: Push image to GitHub Container Registry
         run: |
           echo "check available disk space"
           du -hs
+          cd tts/F5-TTS && docker build . --file Dockerfile --tag $IMAGE_NAME --label "runnumber=${GITHUB_RUN_ID}"
+
+      - name: Push image to GitHub Container Registry
+        run: |
 
           IMAGE_ID=ghcr.io/${{ github.repository_owner }}/$IMAGE_NAME
 
diff --git a/tts/F5-TTS/Dockerfile b/tts/F5-TTS/Dockerfile
index 4c07faf..2a0a431 100644
--- a/tts/F5-TTS/Dockerfile
+++ b/tts/F5-TTS/Dockerfile
@@ -1,5 +1,5 @@
 # Base CUDA image
-FROM cnstark/pytorch:2.3.1-py3.10.15-cuda12.1.0-devel-ubuntu22.04
+FROM pytorch/pytorch:2.5.0-cuda12.1-cudnn9-runtime
 
 # Set environment variables
 ENV DEBIAN_FRONTEND=noninteractive