From f92fffa2dfa242801a8656bf15190c34030d1718 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 10 Nov 2024 17:40:52 +0200
Subject: [PATCH] better warning for missing cuda-toolkit

---
 gimmvfi/generalizable_INR/modules/softsplat.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/gimmvfi/generalizable_INR/modules/softsplat.py b/gimmvfi/generalizable_INR/modules/softsplat.py
index 10ac0aa..369b99b 100644
--- a/gimmvfi/generalizable_INR/modules/softsplat.py
+++ b/gimmvfi/generalizable_INR/modules/softsplat.py
@@ -262,8 +262,11 @@ def cuda_kernel(strFunction: str, strKernel: str, objVariables: typing.Dict):
 
 @cupy.memoize(for_each_device=True)
 def cuda_launch(strKey: str):
-    if "CUDA_HOME" not in os.environ:
-        os.environ["CUDA_HOME"] = cupy.cuda.get_cuda_path()
+    try:
+        os.environ.setdefault("CUDA_HOME", cupy.cuda.get_cuda_path())
+    except Exception:
+        if "CUDA_HOME" not in os.environ:
+            raise RuntimeError("'CUDA_HOME' not set, unable to find cuda-toolkit installation.")
     
     strKernel = objCudacache[strKey]["strKernel"]
     strFunction = objCudacache[strKey]["strFunction"]