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