set_visible_device_list

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

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/source/api_cc/src/DataModifierTF.cc b/source/api_cc/src/DataModifierTF.cc
@@ -52,8 +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.mutable_gpu_options()->visible_device_list =
-        std::to_string(gpu_rank % gpu_num);
+    options.config.mutable_gpu_options()->set_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,8 +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.mutable_gpu_options()->visible_device_list =
-        std::to_string(gpu_rank % gpu_num);
+    options.config.mutable_gpu_options()->set_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,8 +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.mutable_gpu_options()->visible_device_list =
-        std::to_string(gpu_rank % gpu_num);
+    options.config.mutable_gpu_options()->set_visible_device_list(
+        std::to_string(gpu_rank % gpu_num));
     graph::SetDefaultDevice(str, graph_def);
   }
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM