Skip to content

Commit db5f8a0

Browse files
authored
Fix compile with nvjpeg on Windows CUDA 12 (#8641)
1 parent 00e7fa1 commit db5f8a0

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

setup.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -366,7 +366,7 @@ def make_image_extension():
366366
else:
367367
warnings.warn("Building torchvision without AVIF support")
368368

369-
if USE_NVJPEG and torch.cuda.is_available():
369+
if USE_NVJPEG and (torch.cuda.is_available() or FORCE_CUDA):
370370
nvjpeg_found = CUDA_HOME is not None and (Path(CUDA_HOME) / "include/nvjpeg.h").exists()
371371

372372
if nvjpeg_found:
@@ -376,6 +376,8 @@ def make_image_extension():
376376
Extension = CUDAExtension
377377
else:
378378
warnings.warn("Building torchvision without NVJPEG support")
379+
elif USE_NVJPEG:
380+
warnings.warn("Building torchvision without NVJPEG support")
379381

380382
return Extension(
381383
name="torchvision.image",

0 commit comments

Comments
 (0)