name: ium_s444409

docker_env:
  image: s444409-create-dataset
  volumes: ["/tmp/mlruns:/mlruns:rw"]

entry_points:
  main:
    parameters:
      epochs: {type: float, default: 5}
      batch_size: {type: float, default: 64}
    command: "python train_model.py -e {epochs} -b {batch_size}"
  eval:
    command: "python eval_model.py"