BingqingCheng / cace

MIT License
49 stars 10 forks source link

Torchscript #6

Open djkim-kr opened 1 month ago

djkim-kr commented 1 month ago

Now the trained potential is 'Torchscriptable'

Usage: Training procedure remains the same. After training, add the lines to check if there's anything wrong and try to save.

try:
    logging.info("Scripting the model with TorchScript")
    scripted_model = torch.jit.script(cace_nnp)
    print(f"Successfully scripted full_model")
    scripted_model.save("water-model-scripted.pt")
    logging.info("Model successfully scripted and saved as 'water-model-scripted.pt'")
except Exception as e:
    logging.error(f"Error scripting the model: {e}")
    logging.error(traceback.format_exc())

If there's any problem with torchscripting, please inform me immediately

BingqingCheng commented 1 month ago

I think you added all the temporary files in pycache. Please remove and clean up.