version: 1.1
kind: operation
matrix:
  kind: grid
  params:
    sota:
      kind: choice
      value:
        - model: ModelA
          max_epochs: 200
          dataset_hash: ae43ff
        - model: ModelB
          max_epochs: 150
          dataset_hash: 33fba2
component:
  name: model-trainer
  tags:
    - dataset_update_retrain
    - "{{ sota.model }}"  #<- I want to set this dynamically based on the job input
  inputs:
    - name: sota
      type: dict
      isOptional: false
  run:
    kind: job
    container:
      args: >
        set -x;
       train \
          hydra.run.dir={{ globals.run_outputs_path }} \
          model={{ sota.model }} \
            model.datamodule.dataset_hash={{ sota.dataset_hash }}/ \
            model.datamodule.data_directory=/data/{{ sota.model }}/ \
          trainer=gpu \
            trainer.max_epochs={{ sota.max_epochs }} \
      ...