cafbafcfa19b93e4b85d1a6a523c274b7641d3df,main.py,,save_checkpoint,#,244

Before Change

    save_states = {
        "epoch": epoch,
        "arch": arch,
        "state_dict": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict()
    }

After Change

    if hasattr(model, "module"):
        model_state_dict = model.module.state_dict()
    else:
        model_state_dict = model.state_dict()
    save_states = {
        "epoch": epoch,
        "arch": arch,
        "state_dict": model_state_dict,

Instances