From ae06b045dffc173790b3067aa10c9f948e435289 Mon Sep 17 00:00:00 2001
From: Michal Zientkiewicz <michalz@nvidia.com>
Date: Wed, 17 Jul 2024 17:32:15 +0200
Subject: [PATCH] Restrict use of pinned when there's no GPU.

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
---
 dali/pipeline/data/buffer.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/dali/pipeline/data/buffer.cc b/dali/pipeline/data/buffer.cc
index 9097d2ed93d..11b1032dc68 100644
--- 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);
   }();