From 39f357c33f8692b78ab8d6af68ecd73f24916534 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 12 Nov 2024 23:34:07 -0500 Subject: [PATCH] call TFE_ContextOptionsSetConfig Signed-off-by: Jinzhe Zeng --- source/api_cc/src/DeepPotJAX.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/source/api_cc/src/DeepPotJAX.cc b/source/api_cc/src/DeepPotJAX.cc index ac404d5a68..2e49b6e4c1 100644 --- a/source/api_cc/src/DeepPotJAX.cc +++ b/source/api_cc/src/DeepPotJAX.cc @@ -255,6 +255,7 @@ void deepmd::DeepPotJAX::init(const std::string& model, } TF_SetConfig(sessionopts, config.data(), config.size(), status); + check_status(status); TF_Buffer* runopts = NULL; @@ -273,6 +274,8 @@ void deepmd::DeepPotJAX::init(const std::string& model, check_status(status); ctx_opts = TFE_NewContextOptions(); + TFE_ContextOptionsSetConfig(ctx_opts, config.data(), config.size(), status); + check_status(status); ctx = TFE_NewContext(ctx_opts, status); check_status(status); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM