mnist_example.yaml

apiVersion: "kubeflow.org/v1"
kind: "PyTorchJob"
metadata:
  name: "pytorch-dist-mnist-gloo"
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: quay.io/vpavlin/pytorch-dist-mnist-test:1.0-3
              args: ["--backend", "gloo"]
              resources: 
                limits:
                  memory: 1Gi
                  cpu: 1
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers: 
            - name: pytorch
              image: quay.io/vpavlin/pytorch-dist-mnist-test:1.0-3
              args: ["--backend", "gloo"]
              resources: 
                limits:
                  memory: 1Gi
                  cpu: 1