diff --git a/setup.py b/setup.py
index 5c6c6fd7..0b6bdcfa 100644
--- a/setup.py
+++ b/setup.py
@@ -9,9 +9,9 @@
 
 
 # If WITH_CUDA is defined
-if os.environ.get("WITH_CUDA", "0") == "1":
+if os.environ.get("WITH_CUDA") == "1":
     use_cuda = True
-elif os.environ.get("WITH_CUDA", "0") == "0":
+elif os.environ.get("WITH_CUDA") == "0":
     use_cuda = False
 else:
     use_cuda = torch.cuda._is_compiled()