Skip to content

Commit cb7a0d3

Browse files
fix(pt): use eval mode in the C++ interface (#4402)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced error handling for improved clarity during execution. - Added evaluation mode setting for models to ensure correct inference behavior. - **Improvements** - Streamlined device management in computation methods for better resource utilization. - Improved flexibility in handling tensor creation and data transfer. - Updated computation methods to support new logic for message passing and mapping tensors. - Introduced new methods for mixed-type computations, enhancing functionality. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 5d589da commit cb7a0d3

File tree

2 files changed

+2
-0
lines changed

2 files changed

+2
-0
lines changed

source/api_cc/src/DeepPotPT.cc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ void DeepPotPT::init(const std::string& model,
8888
}
8989
std::unordered_map<std::string, std::string> metadata = {{"type", ""}};
9090
module = torch::jit::load(model, device, metadata);
91+
module.eval();
9192
do_message_passing = module.run_method("has_message_passing").toBool();
9293
torch::jit::FusionStrategy strategy;
9394
strategy = {{torch::jit::FusionBehavior::DYNAMIC, 10}};

source/api_cc/src/DeepSpinPT.cc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ void DeepSpinPT::init(const std::string& model,
8888
}
8989
std::unordered_map<std::string, std::string> metadata = {{"type", ""}};
9090
module = torch::jit::load(model, device, metadata);
91+
module.eval();
9192
do_message_passing = module.run_method("has_message_passing").toBool();
9293
torch::jit::FusionStrategy strategy;
9394
strategy = {{torch::jit::FusionBehavior::DYNAMIC, 10}};

0 commit comments

Comments
 (0)