apiVersion: v1
kind: Pod
metadata:
  name: gke-training-pod
spec:
  containers:
  - name: gpu-trainer
    image: gcr.io/deeplearning-platform-release/pytorch-gpu.1-8
    env:
    - name: MODEL_DIR
      value: gs://model_device
    resources:
      limits:
        nvidia.com/gpu: 1
    command: ["python"]
    args: ["trainer/train.py", "--model_dir", "$(MODEL_DIR)"]