From d54856831e0f96727ae01c0a18683257c4e5b78d Mon Sep 17 00:00:00 2001
From: Attila Afra <attila.t.afra@intel.com>
Date: Thu, 14 Mar 2024 20:46:00 +0200
Subject: [PATCH] CUDA: CURTN: fix internal error in cudaGetDevice

---
 devices/cuda/curtn.cpp | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/devices/cuda/curtn.cpp b/devices/cuda/curtn.cpp
index f78fc824..9f300594 100644
--- a/devices/cuda/curtn.cpp
+++ b/devices/cuda/curtn.cpp
@@ -492,14 +492,21 @@ namespace curtn
       if (device == nullptr)
         return Runtime::setError(cudaErrorInvalidValue);
 
-      CUdevice deviceHandle;
-      CUresult result = cuCtxGetDevice(&deviceHandle);
+      CUcontext context;
+      CUresult result = cuCtxGetCurrent(&context);
+      if (result != CUDA_SUCCESS)
+        return Runtime::setError(result);
 
-      if (result == CUDA_SUCCESS)
+      if (context)
       {
+        CUdevice deviceHandle;
+        result = cuCtxGetDevice(&deviceHandle);
+        if (result != CUDA_SUCCESS)
+          return Runtime::setError(result);
+
         result = rt.getDeviceOrdinal(deviceHandle, *device);
       }
-      else if (result == CUDA_ERROR_INVALID_CONTEXT)
+      else
       {
         // No current context, but it will be created lazily for device 0
         *device = 0;