From db5f8a00f74a39233ee00d7549e837c6b205b9de Mon Sep 17 00:00:00 2001 From: Andrey Talman Date: Wed, 11 Sep 2024 12:27:26 -0400 Subject: [PATCH] Fix compile with nvjpeg on Windows CUDA 12 (#8641) --- setup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index dbe8ce58aa2..4b0525d8e41 100644 --- a/setup.py +++ b/setup.py @@ -366,7 +366,7 @@ def make_image_extension(): else: warnings.warn("Building torchvision without AVIF support") - if USE_NVJPEG and torch.cuda.is_available(): + if USE_NVJPEG and (torch.cuda.is_available() or FORCE_CUDA): nvjpeg_found = CUDA_HOME is not None and (Path(CUDA_HOME) / "include/nvjpeg.h").exists() if nvjpeg_found: @@ -376,6 +376,8 @@ def make_image_extension(): Extension = CUDAExtension else: warnings.warn("Building torchvision without NVJPEG support") + elif USE_NVJPEG: + warnings.warn("Building torchvision without NVJPEG support") return Extension( name="torchvision.image",