correctly check cuda initialization in create_cuda_context().

1 files changed, 12 insertions, 8 deletions

diff --git a/gpu.py b/gpu.py
index 78eeb6b..b98d42b 100644
--- a/gpu.py
+++ b/gpu.py
@@ -778,16 +778,20 @@ class GPUPDF(object):
 def create_cuda_context(device_id=None):
     """Initialize and return a CUDA context on the specified device.
     If device_id is None, the default device is used."""
-    try:
-        cuda.mem_get_info()
-    except cuda.LogicError:
-        # initialize cuda
-        cuda.init()
-
     if device_id is None:
-        context = pycuda.tools.make_default_context()
+        try:
+            context = pycuda.tools.make_default_context()
+        except cuda.LogicError:
+            # initialize cuda
+            cuda.init()
+            context = pycuda.tools.make_default_context()
     else:
-        device = cuda.Device(device_id)
+        try:
+            device = cuda.Device(device_id)
+        except cuda.LogicError:
+            # initialize cuda
+            cuda.init()
+            device = cuda.Device(device_id)
         context = device.make_context()
 
     context.set_cache_config(cuda.func_cache.PREFER_L1)