diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml
index 22234a1dc8..5076aebc6f 100644
--- a/docker/docker-compose.yml
+++ b/docker/docker-compose.yml
@@ -34,7 +34,6 @@ services:
   immich-machine-learning:
     container_name: immich_machine_learning
     image: altran1502/immich-machine-learning:release
-    command: [ "python", "src/main.py" ]
     volumes:
       - ${UPLOAD_LOCATION}:/usr/src/app/upload
       - model-cache:/cache
diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile
index ac646a9e0c..acf92053ba 100644
--- a/machine-learning/Dockerfile
+++ b/machine-learning/Dockerfile
@@ -14,3 +14,5 @@ RUN pip install transformers tqdm numpy scikit-learn scipy nltk sentencepiece fl
 RUN pip install --no-deps sentence-transformers
 
 COPY . .
+
+CMD ["python", "src/main.py"]