Init tensor with device agnostic code without it showing up in the state dict

[Robinysh]

This is the current recommended way for initializing tensors inside LightningModules with device-agnostic code.

class LitModel(LightningModule):

    def __init__(self):
        ...
        self.register_buffer("sigma", torch.eye(3))
        # you can now access self.sigma anywhere in your module

Unfortunately, for my current use case, the tensor now shows up in the model state dict as an undesired side effect. My pipeline consists of loading the save file across runs, and sometimes some particular tensors change sizes, which makes my savefile incompatible. I don't want to save and load those particular tensors. Is there a way to work around this?

[tchaton]

Dear @Robinysh,

You could override the state_dict function of your LightningModule and drop those tensors.

Best,
T.C