Error when saving parametrized module

[Linlangxie]

When saving a parametrized module, for my case, adding a torch.nn.utils.parametrize.register_parametrization() to the layers, it raises error: RuntimeError: Serialization of parametrized modules is only supported through state_dict(), at /mase/machop/chop/passes/graph/interface/save_and_load.py line 13, save_graph_module_ckpt().
It seems like torch.save() cannot support the parametrized module to save the model structure.

[pgimenes]

@ChengZhang-98 were these changes pushed to main?