How to I load part of a model weights?

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:
            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?