Update kubeflow sample pytorchjob to pytorch 2.x and cuda 12.x (#1910) (

#2019) Signed-off-by: wangdepeng <wangdepeng_yewu@cmss.chinamobile.com> Co-authored-by: wangdepeng <wangdepeng_yewu@cmss.chinamobile.com>
kubernetes-sigs · Apr 19, 2024 · d4540e3 · d4540e3
1 parent d5968c4
commit d4540e3
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/site/static/examples/jobs/sample-pytorchjob.yaml b/site/static/examples/jobs/sample-pytorchjob.yaml
@@ -14,7 +14,9 @@ spec:
         spec:
           containers:
             - name: pytorch
-              image: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
+              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:v1beta1-21320b6
+#              If you have gpu, pytorch-mnist-gpu would be helpful. pytorch-mnist-gpu is approximately 22GB
+#              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:latest
               imagePullPolicy: Always
               command:
                 - "python3"
@@ -31,7 +33,9 @@ spec:
         spec:
           containers:
             - name: pytorch
-              image: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
+              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:v1beta1-21320b6
+#              If you have gpu, pytorch-mnist-gpu would be helpful. pytorch-mnist-gpu is approximately 22GB
+#              image: docker.io/kubeflowkatib/pytorch-mnist-cpu:latest
               imagePullPolicy: Always
               command:
                 - "python3"