Fix compile with nvjpeg on Windows CUDA 12 (#8641)

pytorch · Sep 11, 2024 · db5f8a0 · db5f8a0
1 parent 00e7fa1
commit db5f8a0
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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",