diff --git a/extensions/cuda_extension.py b/extensions/cuda_extension.py
index edfbf6f24e06..da15bcd57a26 100644
--- a/extensions/cuda_extension.py
+++ b/extensions/cuda_extension.py
@@ -30,7 +30,7 @@ def is_available(self) -> bool:
 
             # 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 b105c03b717c..d2cfb13a4dfd 100644
--- a/setup.py
+++ b/setup.py
@@ -144,6 +144,7 @@ def get_version() -> str:
     package_data={
         "colossalai": [
             "kernel/extensions/csrc/**/*",
+            "kernel/extensions/pybind/**/*",
         ]
     },
 )