Hello,
i have made modifications to a model, expanding it with multiple attention mechanisms. 80% of the model is the same as the previous version. Is it possible, to take the pretrained weights of the previous version, and insert them where applicable?
Ive trained the new model for 1 epoch, saving the weights (checkpoint). This is my attempt at updating those weights with pretrained weights:
    checkpoint = torch.load(args.weight_path)
    pretrained = torch.load("pretrained_weights.pth")
    
    pretrained = {k: v for k, v in pretrained.items() if k in checkpoint}
    checkpoint.update(pretrained)
    
    G_mod.load_state_dict(checkpoint['state_dict'])