name: pytorch_tutorial

python_env: python_env.yaml

entry_points:
  main:
    parameters:
      batch-size: {type: int, default: 64}
      test-batch-size: {type: int, default: 1000}
      epochs: {type: int, default: 10}
      lr: {type: float, default: 0.01}
      momentum: {type: float, default: 0.5}
      enable-cuda: {type: string, default: 'True'}
      seed: {type: int, default: 5}
      log-interval: {type: int, default: 100}
    command: |
          python mnist_tensorboard_artifact.py \
            --batch-size {batch-size} \
            --test-batch-size {test-batch-size} \
            --epochs {epochs} \
            --lr {lr} \
            --momentum {momentum} \
            --enable-cuda {enable-cuda} \
            --seed {seed} \
            --log-interval {log-interval}