workerPoolSpecs:
  - machineSpec:
      machineType: n1-standard-4
      acceleratorType: NVIDIA_TESLA_T4
      acceleratorCount: 1
  - machineSpec:
      machineType: n1-standard-4
      acceleratorType: NVIDIA_TESLA_T4
      acceleratorCount: 1
  replicaCount: 1
  pythonPackageSpec:
    executorImageUri: us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-4:latest
    packageUris: gs://{BUCKET}/flowers-1.0.tar.gz
    pythonModule: flowers.classifier.train
    args:
    - --job-dir="gs://{BUCKET}/flowers_trained/gpus_multiple_machines"  # TODO: Add your bucket!
    - --pattern="-*"
    - --num_epochs=20
    - --distribute="gpus_multiple_machines"