From 718e4018adb5db66004495158e418c5efe94a79f Mon Sep 17 00:00:00 2001 From: Shehtab Zaman Date: Fri, 14 Feb 2025 02:48:28 -0500 Subject: [PATCH] Update setup.py (#360) Check if environment variable is set before retrieving variable at setup time. --- setup.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index 5c6c6fd7..6cd4eee4 100644 --- a/setup.py +++ b/setup.py @@ -9,10 +9,13 @@ # If WITH_CUDA is defined -if os.environ.get("WITH_CUDA", "0") == "1": - use_cuda = True -elif os.environ.get("WITH_CUDA", "0") == "0": - use_cuda = False +if os.getenv("WITH_CUDA") is None: + if os.environ.get("WITH_CUDA", "0") == "1": + use_cuda = True + elif os.environ.get("WITH_CUDA", "0") == "0": + use_cuda = False + else: + raise ValueError("Invalid flag with WITH_CUDA environment variable. Expected '0' or '1'") else: use_cuda = torch.cuda._is_compiled()