From a1e39f4c0debb66d2006a12b0a653999fe570842 Mon Sep 17 00:00:00 2001
From: flybird11111 <1829166702@qq.com>
Date: Thu, 6 Jun 2024 11:47:48 +0800
Subject: [PATCH] [install]fix setup (#5786)
* fix
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
extensions/cuda_extension.py | 2 +-
setup.py | 1 +
2 files changed, 2 insertions(+), 1 deletion(-)
diff --git a/extensions/cuda_extension.py b/extensions/cuda_extension.py
index edfbf6f24..da15bcd57 100644
--- a/extensions/cuda_extension.py
+++ b/extensions/cuda_extension.py
@@ -30,7 +30,7 @@ class _CudaExtension(_CppExtension):
# 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))
+ cuda_available = torch.cuda.is_available() or bool(os.environ.get("FORCE_CUDA", 0))
except:
cuda_available = False
return cuda_available
diff --git a/setup.py b/setup.py
index b105c03b7..d2cfb13a4 100644
--- a/setup.py
+++ b/setup.py
@@ -144,6 +144,7 @@ setup(
package_data={
"colossalai": [
"kernel/extensions/csrc/**/*",
+ "kernel/extensions/pybind/**/*",
]
},
)