restore gpu information printing code

deepmodeling · Jan 8, 2025 · a749056 · a749056
1 parent 9f558e7
commit a749056
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/source/api_cc/src/DeepPotPD.cc b/source/api_cc/src/DeepPotPD.cc
@@ -106,17 +106,18 @@ void DeepPotPD::init(const std::string& model,
     throw deepmd::deepmd_exception("Given inference model: " + model +
                                    " do not exist, please check it.");
   }
-
   const char* use_cuda_toolkit = std::getenv("USE_CUDA_TOOLKIT");
   gpu_enabled = (use_cuda_toolkit && (std::string(use_cuda_toolkit) == "1"));
   config->SetModel(pdmodel_path, pdiparams_path);
   config_fl->SetModel(fl_pdmodel_path, fl_pdiparams_path);
   if (!gpu_enabled) {
     config->DisableGpu();
     config_fl->DisableGpu();
+    std::cout << "load model from: " << model << " to cpu " << std::endl;
   } else {
     config->EnableUseGpu(4096, 0);
     config_fl->EnableUseGpu(4096, 0);
+    std::cout << "load model from: " << model << " to gpu:" << gpu_id << std::endl;
   }
 
   // NOTE: Both set to 1 now.