How to save modelcheckpoint when use spectral_norm

[ZengPan-github]

When using spectral_normalization in ddp models, the following error occurs

RuntimeError: Serialization of parametrized modules is only supported through state_dict(). See:
https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training

Is there a proper way to solve this problem?

Reproduction script

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import ModelCheckpoint
from torch.nn.utils.parametrizations import spectral_norm
import torch.nn as nn

class net(LightningModule):
    def __init__(self):
        super(net, self).__init__()
        self.l1 = nn.Sequential(
            spectral_norm(nn.Conv2d(1, 20, 5)),
            nn.ReLU(),
            spectral_norm(nn.Conv2d(20, 64, 5)),
            nn.ReLU()
        )

    def forward(self, x):
        y = self.l1(x)
        return y

    def training_step(self, *args, **kwargs):
        ...


model = net()
modelckpt = ModelCheckpoint(dirpath='path', filename='name', monitor='val_loss', save_top_k=3)
trainer = Trainer(accelerator='gpu', devices='0,1', callbacks=[modelckpt])
trainer.fit(model)