Skip to content

Commit 1707266

Browse files
committed
fix pack .h in whl problem
1 parent ecc1653 commit 1707266

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

setup.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,11 @@ def build_cmake(self, ext: CMakeExtension):
7676
cmake_args += ["-DRUNTIME_ENVIRONMENT=maca"]
7777
cmake_args += ["-DBUILD_UCM_SPARSE=OFF"]
7878
case _:
79-
# default ascend
80-
cmake_args += ["-DRUNTIME_ENVIRONMENT=ascend"]
79+
import torch
80+
if hasattr(torch, "cuda") and torch.cuda.is_available():
81+
cmake_args += ["-DRUNTIME_ENVIRONMENT=cuda"]
82+
else:
83+
cmake_args += ["-DRUNTIME_ENVIRONMENT=ascend"]
8184

8285
subprocess.check_call(
8386
["cmake", *cmake_args, ext.cmake_file_path], cwd=build_dir
@@ -103,4 +106,5 @@ def build_cmake(self, ext: CMakeExtension):
103106
ext_modules=[CMakeExtension(name="ucm", source_dir=ROOT_DIR)],
104107
cmdclass={"build_ext": CMakeBuild},
105108
zip_safe=False,
109+
include_package_data=False,
106110
)

0 commit comments

Comments
 (0)