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/**/*",
         ]
     },
 )