name: s487194

docker_env:
  image: s487194/ium:1

entry_points:
  train:
    parameters:
      epochs: {type: int, default: 100}
      lr: {type: float, default: 0.01}
    command: "python train.py --epochs {epochs} --lr {lr}"
    
  test:
    command: "python prediction.py"