import  optimizer_config

training_loop.Trainer:
  learning_rate_schedule = @optimizer_config.lr_cosine_decay

optimizer_config.lr_cosine_decay:
    max_lr = 0.01
    min_lr = 0.001
    decay_after = True
    spike_steps = 0
    spike_lr = 0.0