From 2d6992d4236ce74d6a8af4247735a30a21ecdbf7 Mon Sep 17 00:00:00 2001 From: Charles Coulombe Date: Thu, 28 Mar 2024 10:08:35 -0400 Subject: [PATCH] Allow building cuda extension without a device. 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 f1e0095b29b6..c7fd0df16b22 100644 --- a/extensions/cuda_extension.py +++ b/extensions/cuda_extension.py @@ -27,7 +27,9 @@ def is_available(self) -> bool: 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