diff --git a/configs/volcano-trainjob-integration.yaml b/configs/volcano-trainjob-integration.yaml
index 77fb7a9..0b14d7b 100644
--- a/configs/volcano-trainjob-integration.yaml
+++ b/configs/volcano-trainjob-integration.yaml
@@ -310,6 +310,9 @@ spec:
       replicatedJobs:
       - name: node
         template:
+          metadata:
+            labels:
+              trainer.kubeflow.org/trainjob-ancestor-step: trainer
           spec:
             template:
               metadata:
@@ -376,7 +379,7 @@ spec:
   trainer:
     image: nvcr.io/nvidia/pytorch:24.10-py3
     command:
-      - python
+      - torchrun
       - /workspace/scripts/train_nccl.py
     numNodes: 2
     resourcesPerNode: