File tree Expand file tree Collapse file tree 1 file changed +6
-2
lines changed
Expand file tree Collapse file tree 1 file changed +6
-2
lines changed Original file line number Diff line number Diff 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)
You can’t perform that action at this time.
0 commit comments