From c46e09715cc906b9c82a92f1287bd1b1defe87ea Mon Sep 17 00:00:00 2001 From: Charles Coulombe Date: Wed, 5 Jun 2024 02:26:30 -0400 Subject: [PATCH] Allow building cuda extension without a device. (#5535) Added FORCE_CUDA environment variable support, to enable building extensions where a GPU device is not present but cuda libraries are. --- extensions/cuda_extension.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/extensions/cuda_extension.py b/extensions/cuda_extension.py index b722057c9..edfbf6f24 100644 --- a/extensions/cuda_extension.py +++ b/extensions/cuda_extension.py @@ -28,7 +28,9 @@ class _CudaExtension(_CppExtension): try: import torch - cuda_available = torch.cuda.is_available() + # torch.cuda.is_available requires a device to exist, allow building with cuda extension on build nodes without a device + # but where cuda is actually available. + cuda_available = torch.cuda.is_available() or bool(os.environ.get('FORCE_CUDA', 0)) except: cuda_available = False return cuda_available