diff --git a/Dockerfile b/Dockerfile
index 4cc4308..713d33b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -94,7 +94,7 @@ ENV DISABLE_MODEL_DOWNLOAD=false
 ENV DISABLE_TRAINING_ASSET_DOWNLOAD=false
 ENV SHUTDOWN_AFTER_PROVISION=false
 ENV DISABLE_AUTOLAUNCH=false
-ENV POD_TERMINATION_PID_PATTERN=train
+ENV POD_TERMINATION_PID_PATTERN=train.py
 
 # Copy the scripts
 COPY --chmod=755 scripts/* ./