Get and run submodule from a deployed model

[Hongyu-yu]

Hello, I wonder if I can get and run submodule from a deployed model.

To be more explicit, I have a model trained with StressForceOutput and then deployed as jit.pth. Then I load the jit.pth with ase interface NequIPCalculator.from_deployed_model. Now I want to use the unwrapped energy model func in StressForceOutput from jit.pth so that I can forward the model without grad wrapper.

I try to use get_submodule() and other method but return nothing but barely RecursiveScriptModule(original_name=RescaleOutput).
What should I do? Many thanks!

[Linux-cpp-lisp]

Hm GradientOutput has a skip flag (https://github.com/mir-group/nequip/blob/main/nequip/nn/_grad_output.py#L26) for exactly this purpose, but StressForceOutput is missing one.

Probably the easiest thing to do is to save the state_dict() of the deployed model (which should have correct key names still, I think?) and then create a new model without the StressForceOutput model builder , and then use load_model_state builder to load the weights from the deployed model. You can deploy this YAML directly with nequip-deploy build --model

[Hongyu-yu]

Got it! Thanks!