From 530f98e995af66b71faefe5eba695185b27c7449 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 22 Nov 2024 14:18:23 -0500 Subject: [PATCH] Revert "Freeze TorchScript modules with preserved attributes" This reverts commit e4453720c71d80af0f92e196ab1bf9f69b4b254e. --- source/api_cc/src/DeepPotPT.cc | 2 -- source/api_cc/src/DeepSpinPT.cc | 2 -- 2 files changed, 4 deletions(-) diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 4f17c30e25..b431ad65cf 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -89,8 +89,6 @@ void DeepPotPT::init(const std::string& model, std::unordered_map metadata = {{"type", ""}}; module = torch::jit::load(model, device, metadata); module.eval(); - const std::vector& preserved_attrs = {"forward_lower", "has_message_passing", "get_rcut", "get_ntypes", "get_dim_fparam", "get_dim_aparam", "is_aparam_nall", "get_type_map"}; - module = torch::jit::freeze(module, preserved_attrs); do_message_passing = module.run_method("has_message_passing").toBool(); torch::jit::FusionStrategy strategy; strategy = {{torch::jit::FusionBehavior::DYNAMIC, 10}}; diff --git a/source/api_cc/src/DeepSpinPT.cc b/source/api_cc/src/DeepSpinPT.cc index fac494772c..1a245c7b2e 100644 --- a/source/api_cc/src/DeepSpinPT.cc +++ b/source/api_cc/src/DeepSpinPT.cc @@ -89,8 +89,6 @@ void DeepSpinPT::init(const std::string& model, std::unordered_map metadata = {{"type", ""}}; module = torch::jit::load(model, device, metadata); module.eval(); - const std::vector& preserved_attrs = {"forward_lower", "has_message_passing", "get_rcut", "get_ntypes", "get_dim_fparam", "get_dim_aparam", "is_aparam_nall", "get_type_map"}; - module = torch::jit::freeze(module, preserved_attrs); do_message_passing = module.run_method("has_message_passing").toBool(); torch::jit::FusionStrategy strategy; strategy = {{torch::jit::FusionBehavior::DYNAMIC, 10}};