I am not sure how model and optimizers work together in pytorch.
Here is the thing,
I save my model(full model, not the state_dict of the model) and i save the optimizer.state_dict
torch.save({
            'model': model, # it saves the whole model
            'optimizer_state_dict': optimizer.state_dict(),
            'lr_scheduler_state_dict': lr_scheduler.state_dict(),
            }, save_path)

Then i load my model and freeze some layers and i define the optimizer again and then load the optimizers state_dict…
ckpt = torch.load(save_path) 
model = ckpt['model']

for name, param in model.named_parameters():
    if ('layer4' in name) or ('fc' in name):
        param.requires_grad = True 
    else:
        param.requires_grad = False
        
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = lr)
optimizer.load_state_dict(ckpt['optimizer_state_dict'])

exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
exp_lr_scheduler.load_state_dict(ckpt['lr_scheduler'])

It throws an ugly error-
ValueError: loaded state dict contains a parameter group that doesn’t match the size of optimizer’s group
Can you pls help me get it corrected ?
Any why do we even need to save the state_dict of the optimizers and the scheduler ?