diff --git a/llm/gpt-2/gpt2-train.yaml b/llm/gpt-2/gpt2-train.yaml
index 7c79f53ee44..7c7e58eeaf9 100644
--- a/llm/gpt-2/gpt2-train.yaml
+++ b/llm/gpt-2/gpt2-train.yaml
@@ -7,7 +7,7 @@ resources:
   accelerators: A100:8
   # Use docker image for latest version g++ to enable the compilation of llm.c.
   image_id: docker:nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
-  image_id:
+  any_of:
     # Avoid using docker image for lambda due to the docker is not supported on
     # Lambda yet, but the base image works.
     - cloud: lambda
@@ -16,11 +16,12 @@ resources:
     - cloud: gcp
     - cloud: azure
     - cloud: fluidstack
+    - cloud: kubernetes
   
 file_mounts:
-  ~/.cache/huggingface: gs://$BUCKET_NAME
-    # name: $BUCKET_NAME
-    # mode: COPY
+  ~/.cache/huggingface:
+    name: $BUCKET_NAME
+    mode: COPY
 
 setup: |
   cd ~