diff --git a/.github/workflows/training.yml b/.github/workflows/training.yml
index 4ec1ce0..3199858 100644
--- a/.github/workflows/training.yml
+++ b/.github/workflows/training.yml
@@ -3,10 +3,13 @@ on: workflow_dispatch
 
 jobs:
   container-train:
-    runs-on: ubuntu-latest
     name: Container Train
-    container:
-      image: andreaschandra/sunflow:latest
+    runs-on: ubuntu-latest
     steps:
-      - name: Running train script
-        run: python src/train.py
+      - name: Docker run
+        uses: addnab/docker-run-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+          image: andreaschandra/sunflow:latest
+          run: python3 src/train.py
diff --git a/Dockerfile b/Dockerfile
index 19153c6..6817c65 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -26,4 +26,4 @@ RUN pip3 install -r requirements.txt
 
 RUN pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
 
-ENTRYPOINT ["python3", "src/train.py"]
\ No newline at end of file
+CMD ["python3", "src/train.py"]
\ No newline at end of file