code review changes

deepanker13 · Jan 8, 2024 · db7ea06 · db7ea06
1 parent 692e8d8
commit db7ea06
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 28 deletions.
diff --git a/.github/workflows/publish-example-images.yaml b/.github/workflows/publish-example-images.yaml
@@ -52,7 +52,9 @@ jobs:
           - component-name: mxnet-auto-tuning
             dockerfile: examples/mxnet/tune/Dockerfile
             context: examples/mxnet/tune
-
+          - component-name: train-api-hf-image
+            dockerfile: sdk/python/kubeflow/trainer/hf_dockerfile
+            context: sdk/python/kubeflow/trainer
 # TODO (tenzen-y): Fix the below broken Dockerfiles
 #          - component-name: pytorch-dist-mnist-mpi
 #            dockerfile: examples/pytorch/mnist/Dockerfile-mpi

diff --git a/.github/workflows/publish-sdk-images.yaml b/.github/workflows/publish-sdk-images.yaml
diff --git a/sdk/python/kubeflow/trainer/hf_dockerfile b/sdk/python/kubeflow/trainer/hf_dockerfile
@@ -1,5 +1,5 @@
 # Use an official Pytorch runtime as a parent image
-FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+FROM nvcr.io/nvidia/pytorch:23.12-py3
 
  # Set the working directory in the container
  WORKDIR /app