Restrict use of pinned when there's no GPU.

Signed-off-by: Michal Zientkiewicz <[email protected]>
NVIDIA · Jul 17, 2024 · ae06b04 · ae06b04
1 parent ddb2f5c
commit ae06b04
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/dali/pipeline/data/buffer.cc b/dali/pipeline/data/buffer.cc
@@ -64,11 +64,11 @@ DLL_PUBLIC shared_ptr<uint8_t> AllocBuffer(size_t bytes, bool pinned,
 
 DLL_PUBLIC bool RestrictPinnedMemUsage() {
   static bool val = []() {
-    /* int n = 0;
-    if (cudaGetDeviceCount(&n) != CUDA_SUCCESS)
+    int n = 0;
+    if (cudaGetDeviceCount(&n) != cudaSuccess)
       return true;
     if (n == 0)
-      return true; */
+      return true;
     const char *env = getenv("DALI_RESTRICT_PINNED_MEM");
     return env && atoi(env);
   }();