gpu_options

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Sep 30, 2024 · da2e7ec · da2e7ec
1 parent d3944bb
commit da2e7ec
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/source/api_cc/src/DataModifierTF.cc b/source/api_cc/src/DataModifierTF.cc
@@ -52,7 +52,8 @@ void DipoleChargeModifierTF::init(const std::string& model,
     std::string str = "/gpu:0";
     // See
     // https://github.com/tensorflow/tensorflow/blame/8fac27b486939f40bc8e362b94a16a4a8bb51869/tensorflow/core/protobuf/config.proto#L80
-    options.config.visible_device_list = std::to_string(gpu_rank % gpu_num);
+    options.config.mutable_gpu_options()->visible_device_list =
+        std::to_string(gpu_rank % gpu_num);
     graph::SetDefaultDevice(str, graph_def);
   }
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

diff --git a/source/api_cc/src/DeepPotTF.cc b/source/api_cc/src/DeepPotTF.cc
@@ -450,7 +450,8 @@ void DeepPotTF::init(const std::string& model,
     std::string str = "/gpu:0";
     // See
     // https://github.com/tensorflow/tensorflow/blame/8fac27b486939f40bc8e362b94a16a4a8bb51869/tensorflow/core/protobuf/config.proto#L80
-    options.config.visible_device_list = std::to_string(gpu_rank % gpu_num);
+    options.config.mutable_gpu_options()->visible_device_list =
+        std::to_string(gpu_rank % gpu_num);
     graph::SetDefaultDevice(str, graph_def);
   }
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

diff --git a/source/api_cc/src/DeepTensorTF.cc b/source/api_cc/src/DeepTensorTF.cc
@@ -49,7 +49,8 @@ void DeepTensorTF::init(const std::string &model,
     std::string str = "/gpu:0";
     // See
     // https://github.com/tensorflow/tensorflow/blame/8fac27b486939f40bc8e362b94a16a4a8bb51869/tensorflow/core/protobuf/config.proto#L80
-    options.config.visible_device_list = std::to_string(gpu_rank % gpu_num);
+    options.config.mutable_gpu_options()->visible_device_list =
+        std::to_string(gpu_rank % gpu_num);
     graph::SetDefaultDevice(str, graph_def);
   }
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM