Modules with SpectralNormGouk cannot be saved

[Zhangyanbo]

Modules with Gouk's version of spectral norm cannot be saved by torch.save. However, the native SN of PyTorch don't has this problem.

Reproduce:

import torch
import torch.nn as nn

from iResnet import SpectralNormGouk as sn


model = sn.spectral_norm(nn.Conv2d(3, 3, 3))
torch.save(model, 'gouk.pth')

Error:

'''
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/var/folders/r4/n7vnlt1528b_18nsjk7v_02w0000gn/T/ipykernel_81895/1404222307.py in <module>
      1 model = sn.spectral_norm(nn.Conv2d(3, 3, 3))
      2 
----> 3 torch.save(model, 'gouk.pth')

/opt/homebrew/Caskroom/miniforge/base/envs/ML/lib/python3.8/site-packages/torch/serialization.py in save(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization)
    377         if _use_new_zipfile_serialization:
    378             with _open_zipfile_writer(opened_file) as opened_zipfile:
--> 379                 _save(obj, opened_zipfile, pickle_module, pickle_protocol)
    380                 return
    381         _legacy_save(obj, opened_file, pickle_module, pickle_protocol)

/opt/homebrew/Caskroom/miniforge/base/envs/ML/lib/python3.8/site-packages/torch/serialization.py in _save(obj, zip_file, pickle_module, pickle_protocol)
    482     pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol)
    483     pickler.persistent_id = persistent_id
--> 484     pickler.dump(obj)
    485     data_value = data_buf.getvalue()
    486     zip_file.write_record('data.pkl', data_value, len(data_value))

AttributeError: Can't pickle local object 'SpectralNorm.apply.<locals>.<lambda>'
'''