I have modified a model by removing some stuff at the end of it. But of course when I try to load the weights, it complains that certain keys dont exist in the model.
I tried to remove the keys as follows:
            # Remove keys contaning second stage
            keys_to_remove = []
            for key in ckpt_state.keys():
                if "rroi" in key:
                    keys_to_remove.append(key)
            for key in keys_to_remove:
                del ckpt_state[key]


And the models load state dict worked, but the optimizer load_state_dict crashes with
ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group

How to I remove the right stuff from the optimizer weights?