Update setup.py

phuongman · Dec 1, 2024 · 67b0139 · 67b0139
1 parent edb4480
commit 67b0139
Showing 1 changed file with 8 additions and 15 deletions.
diff --git a/networks/channelnorm_package/setup.py b/networks/channelnorm_package/setup.py
@@ -1,27 +1,20 @@
 #!/usr/bin/env python3
 import os
 import torch
-
 from setuptools import setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CppExtension  # Thay CUDAExtension bằng CppExtension
 
-cxx_args = ['-std=c++11']
+cxx_args = ['-std=c++11']  # Chỉ giữ các tham số C++ cần thiết
 
-nvcc_args = [
-    '-gencode', 'arch=compute_52,code=sm_52',
-    '-gencode', 'arch=compute_60,code=sm_60',
-    '-gencode', 'arch=compute_61,code=sm_61',
-    '-gencode', 'arch=compute_70,code=sm_70',
-    '-gencode', 'arch=compute_70,code=compute_70'
-]
+# Xóa nvcc_args vì không sử dụng CUDA
 
 setup(
-    name='channelnorm_cuda',
+    name='channelnorm_cpu',  # Đổi tên để chỉ ra rằng đây là phiên bản chạy trên CPU
     ext_modules=[
-        CUDAExtension('channelnorm_cuda', [
-            #'channelnorm_cuda.cc',
-            'channelnorm_kernel.cu'
-        ], extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args})
+        CppExtension('channelnorm_cpu', [  # Sử dụng CppExtension thay vì CUDAExtension
+            #'channelnorm_cuda.cc',  # Loại bỏ các tệp CUDA
+            'channelnorm_kernel.cpp'  # Chỉ sử dụng tệp C++ thay vì CUDA
+        ], extra_compile_args={'cxx': cxx_args})  # Không cần nvcc_args nữa
     ],
     cmdclass={
         'build_ext': BuildExtension