From 215325c467b3982484abd9f16d45c4919865fa75 Mon Sep 17 00:00:00 2001
From: Julius Richter <julius.richter@uni-hamburg.de>
Date: Wed, 3 Jul 2024 16:10:41 +0200
Subject: [PATCH] skip loading of upfirn2d to run without CUDA

---
 sgmse/backbones/ncsnpp_utils/op/__init__.py  |  3 +--
 sgmse/backbones/ncsnpp_utils/op/upfirdn2d.py | 17 ++++++++++-------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/sgmse/backbones/ncsnpp_utils/op/__init__.py b/sgmse/backbones/ncsnpp_utils/op/__init__.py
index d0918d9..857a9fc 100755
--- a/sgmse/backbones/ncsnpp_utils/op/__init__.py
+++ b/sgmse/backbones/ncsnpp_utils/op/__init__.py
@@ -1,2 +1 @@
-from .fused_act import FusedLeakyReLU, fused_leaky_relu
-from .upfirdn2d import upfirdn2d
+from .upfirdn2d import upfirdn2d
\ No newline at end of file
diff --git a/sgmse/backbones/ncsnpp_utils/op/upfirdn2d.py b/sgmse/backbones/ncsnpp_utils/op/upfirdn2d.py
index a4cf05d..0f4dc56 100755
--- a/sgmse/backbones/ncsnpp_utils/op/upfirdn2d.py
+++ b/sgmse/backbones/ncsnpp_utils/op/upfirdn2d.py
@@ -7,14 +7,17 @@ from torch.utils.cpp_extension import load
 
 
 module_path = os.path.dirname(__file__)
-upfirdn2d_op = load(
-    "upfirdn2d",
-    sources=[
-        os.path.join(module_path, "upfirdn2d.cpp"),
-        os.path.join(module_path, "upfirdn2d_kernel.cu"),
-    ],
-)
 
+if torch.cuda.is_available():
+    upfirdn2d_op = load(
+        "upfirdn2d",
+        sources=[
+            os.path.join(module_path, "upfirdn2d.cpp"),
+            os.path.join(module_path, "upfirdn2d_kernel.cu"),
+        ],
+    )
+else:
+    upfirdn2d_op = None
 
 class UpFirDn2dBackward(Function):
     @staticmethod
-- 
GitLab