wardhaddadin1
commented
6 months ago Hello Peter!
Yeah it's not as simple as just calling torch.jit.script because there are objects attached to the model and module which are not torch scriptable (like losses etc..).
If you want to a get an openmm torchscipt compatible model, you can do (which is used internally in the openmm-ml interface)
from physicsml.plugins.openmm.load import to_openmm_torchscript
ts_model = to_openmm_torchscript(
model_path: Optional[str] = None, # path to saved model
atom_list: Optional[List[int]] = None, # list of atomic numbers of system
system_path: Optional[str] = None, # or path to a system file (like a pdb file)
atom_idxs: Optional[List[int]] = None, # indices of atoms to be used in the torchscript model (for mixed systems)
y_output: Optional[str] = "y_graph_scalars", # the output of the model
pbc: Optional[Tuple[bool, bool, bool]] = None, # pbcs
cell: Optional[List[List[float]]] = None, # cell dims
output_scaling: Optional[float] = None, # whether to scale the output for different units
position_scaling: Optional[float] = None, # whether to scale the input positions (to match training data units)
device: str = 'cpu', # cpu or cuda
precision: str = '32', # 32 or 64
torchscipt_path: Optional[str] = None, # whether to write out the model to a torchscipt file as well
) Let me know if this works for you!
Best, Ward
peastman
I've trained a tensornet model and saved it to disk. I now want to load it in and compile it to torchscript.
It fails with this error.
I also tried compiling just the contained module:
That fails with a different exception:
To work around that problem, I tried creating a new Trainer to attach:
That leads to still a different exception:
This is with Python 3.11.8, PyTorch 2.0.1, and PhysicsML 0.2.0.