seml:
  executable: 'train_seml.py'
  output_dir: '~/logs'
  conda_environment: 'tensorflow'
  project_root_dir: '.'

slurm:
  experiments_per_job: 1
  sbatch_options:
    gres: gpu:1
    mem:  16G
    cpus-per-task: 2
    time: 30-00:00
    partition: gpu_all

fixed:
  model_name: "dimenet++"

  emb_size: 128
  out_emb_size: 256
  int_emb_size: 64
  basis_emb_size: 8
  num_blocks: 4

  num_spherical: 7
  num_radial: 6

  extensive: True
  cutoff: 5.0
  envelope_exponent: 5

  num_before_skip: 1
  num_after_skip: 2
  num_dense_output: 3

  num_train: 110000
  num_valid: 10000
  data_seed: 42
  dataset: "./data/qm9_eV.npz"
  logdir: "/nfs/staff-ssd/klicpera/dimenet/logs"

  num_steps: 3000000
  ema_decay: 0.999

  learning_rate: 0.001
  warmup_steps: 3000
  decay_rate: 0.01
  decay_steps: 4000000

  batch_size: 32
  evaluation_interval: 10000
  save_interval: 10000
  restart: None
  comment: "DimeNet++"

zero_init:
  fixed:
    output_init: 'zeros'
  grid:
    targets:
      type: choice
      options:
        - ['mu']
        - ['homo']
        - ['lumo']
        # - ['gap']
        - ['zpve']

glorot_init:
  fixed:
    output_init: 'GlorotOrthogonal'
  grid:
    targets:
      type: choice
      options:
        - ['alpha']
        - ['r2']
        - ['U0']
        - ['U']
        - ['H']
        - ['G']
        - ['Cv']