# Buggy code
def __init__(self, cfg: DictConfig, fairseq_optimizer):
    super().__init__(cfg, fairseq_optimizer)
    if isinstance(cfg.lr, Collection) and len(cfg.lr) > 1:
        raise ValueError(
            "Cannot use a fixed learning rate schedule with cosine."
            " Consider --lr-scheduler=fixed instead."
        )

    warmup_end_lr = cfg.max_lr
    lr = cfg.lr[0] if isinstance(cfg.lr, Collection) else cfg.lr
    if cfg.warmup_init_lr < 0:
        cfg.warmup_init_lr = lr

    self.min_lr = lr
    self.max_lr = cfg.max_lr
    assert self.max_lr > self.min_lr, "max_lr must be more than lr"

    self.t_mult = cfg.t_mult
    self.period = cfg.lr_period_updates

    if self.period <= 0:
        assert cfg.max_update >= 0, (
            "Either --max_update or --lr-period-updates must be set"
        )
        self.period = cfg.max_update - cfg.warmup_updates

    if cfg.warmup_updates > 0:
        # linearly warmup for the first args.warmup_updates
        self.lr_step = (warmup_end_lr - cfg.warmup_init_lr) / cfg.warmup_updates
    else:
        self.lr_step = 1

    self.warmup_updates = cfg.warmup_updates
    self.lr_shrink = cfg.lr_shrink

    # initial learning rate
    self.lr = cfg.warmup_init_lr
    self.optimizer.set_lr(self.lr)


# Traceback (TypeError)
/experiments/falva/tools/fairseq/fairseq/models/fairseq_model.py:280: UserWarning: FairseqModel is deprecated, please use FairseqEncoderDecoderModel or BaseFairseqModel instead
for key in self.keys
Traceback (most recent call last):
File "/home/falva/anaconda3/envs/mtl4ts/bin/fairseq-generate", line 11, in <module>
load_entry_point('fairseq', 'console_scripts', 'fairseq-generate')()
File "/experiments/falva/tools/fairseq/fairseq_cli/generate.py", line 190, in cli_main
main(args)
File "/experiments/falva/tools/fairseq/fairseq_cli/generate.py", line 47, in main
task=task,
File "/experiments/falva/tools/fairseq/fairseq/checkpoint_utils.py", line 167, in load_model_ensemble
ensemble, args, _task = load_model_ensemble_and_task(filenames, arg_overrides, task)
File "/experiments/falva/tools/fairseq/fairseq/checkpoint_utils.py", line 186, in load_model_ensemble_and_task
model.load_state_dict(state['model'], strict=True, args=args)
TypeError: load_state_dict() got an unexpected keyword argument 'args'

# Fixed code
def __init__(self, cfg: DictConfig, fairseq_optimizer):
    super().__init__(cfg, fairseq_optimizer)
    if isinstance(cfg.lr, Collection) and len(cfg.lr) > 1:
        raise ValueError(
            "Cannot use a fixed learning rate schedule with cosine."
            f" Consider --lr-scheduler=fixed instead. ({cfg.lr})"
        )

    warmup_end_lr = cfg.max_lr
    lr = cfg.lr[0] if isinstance(cfg.lr, Collection) else cfg.lr
    if cfg.warmup_init_lr < 0:
        cfg.warmup_init_lr = lr

    self.min_lr = lr
    self.max_lr = cfg.max_lr
    assert self.max_lr > self.min_lr, "max_lr must be more than lr"

    self.t_mult = cfg.t_mult
    self.period = cfg.lr_period_updates

    if self.period <= 0:
        assert cfg.max_update >= 0, (
            "Either --max_update or --lr-period-updates must be set"
        )
        self.period = cfg.max_update - cfg.warmup_updates

    if cfg.warmup_updates > 0:
        # linearly warmup for the first args.warmup_updates
        self.lr_step = (warmup_end_lr - cfg.warmup_init_lr) / cfg.warmup_updates
    else:
        self.lr_step = 1

    self.warmup_updates = cfg.warmup_updates
    self.lr_shrink = cfg.lr_shrink

    # initial learning rate
    self.lr = cfg.warmup_init_lr
    self.optimizer.set_lr(self.lr)