fix(torch): Raise an exception if gpu is not available

jolibrain · Nov 9, 2020 · 1f0887a · 1f0887a
1 parent eb74a51
commit 1f0887a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/backends/torch/torchlib.cc b/src/backends/torch/torchlib.cc
@@ -141,7 +141,12 @@ namespace dd
     if (lib_ad.has("template"))
       _template = lib_ad.get("template").get<std::string>();
     if (lib_ad.has("gpu"))
-      gpu = lib_ad.get("gpu").get<bool>() && torch::cuda::is_available();
+      gpu = lib_ad.get("gpu").get<bool>();
+    if (gpu && !torch::cuda::is_available())
+      {
+        throw MLLibBadParamException(
+            "GPU is not available, service could not be created");
+      }
     if (lib_ad.has("gpuid"))
       gpuid = lib_ad.get("gpuid").get<int>();
     if (lib_ad.has("nclasses"))