Is there a way to save only part of the Lightning sub-modules to the checkpoint file?

[ahikaml]

I'll explain: Let's say that I have two nn.modules inside my main LightningModule, but one of them is frozen, i.e. doesn't learn during the training but is only used for inferencing during training (requires_grad is False in this module) and I would like to avoid saving the state_dictionray of this static (frozen) module to the checkpoint file.

In plain PyTorch I'd probably filter manually the state_dictionray fields of the frozen module before the saving.
Is there a simple way to do that with pytorch-lightning? Or to raise some flag inside the modules which say to the LightningModule not to save all the parameters inside this frozen module?

A simple toy example for clarification.
In this example, I'd like to avoid saving the parameters of self.frozen_nn_module.
All parameters in self.frozen_nn_module don't require grads.

class LightMod(LightningModule):
    def __init__(self):
        super().__init__()

        #some non frozen module  
        self.non_frozen_nn_module = non_frozen_nn_module
        #some frozen(static) nn.Module
        self.frozen_nn_module= frozen_nn_module

    def forward(self, x):
    Some code....

[rohitgr7]

you have to do that here too.
within lightning you can override on_save_checkpoint hook of LightningModule.

def on_save_checkpoint(self, checkpoint):
    checkpoint['state_dict'] <- remove/pop keys from here