f0144bfcba8411fd68a934158307ec49b00a0ea2,vocoder/train.py,,main,#Any#,430

Before Change

    // if num_gpus > 1:
    //     model = apply_gradient_allreduce(model)

    if c.noam_schedule:
        scheduler_gen = NoamLR(optimizer_gen,
                               warmup_steps=c.warmup_steps_gen,
                               last_epoch=args.restore_step - 1)
        scheduler_disc = NoamLR(optimizer_disc,
                                warmup_steps=c.warmup_steps_gen,
                                last_epoch=args.restore_step - 1)
    else:
        scheduler_gen, scheduler_disc = None, None

    num_params = count_parameters(model_gen)
    print(" > Generator has {} parameters".format(num_params), flush=True)
    num_params = count_parameters(model_disc)
    print(" > Discriminator has {} parameters".format(num_params), flush=True)

After Change

            model_disc.load_state_dict(checkpoint["model_disc"])
            print(" > Restoring Discriminator Optimizer...")
            optimizer_disc.load_state_dict(checkpoint["optimizer_disc"])
            if "scheduler" in checkpoint:
                print(" > Restoring Generator LR Scheduler...")
                scheduler_gen.load_state_dict(checkpoint["scheduler"])
                // NOTE: Not sure if necessary
                scheduler_gen.optimizer = optimizer_gen
            if "scheduler_disc" in checkpoint:
                print(" > Restoring Discriminator LR Scheduler...")
                scheduler_disc.load_state_dict(checkpoint["scheduler_disc"])
                scheduler_disc.optimizer = optimizer_disc

Instances