name: mnist-torchscript

python_env: python_env.yaml

entry_points:
  main:
    parameters:
      epochs: {type: int, default: 5}
      batch_size: {type: int, default: 64}
      learning_rate: {type: float, default: 1e-3}

    command: |
          python mnist_torchscript.py \
            --epochs {epochs} \
            --batch-size {batch_size} \
            --lr {learning_rate}