how should I set it up to train and save the model checkpoint of the module I added?

[youmo445]

Hello, thank you for your work. May I ask, if I want to add a module to WanVideoUnit, such as adding an action embedding module and combining the action embedding and reference image embedding, how should I set it up to train and save the model checkpoint of the module I added?