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