Skip to content

GitLab

Commit ac36d82e authored by Casper's avatar Casper
Browse files

Fix torch import

parent 83c2f09d
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
setup.py
View file @ ac36d82e
...@@ -3,7 +3,7 @@ import torch ...@@ -3,7 +3,7 @@ import torch
from pathlib import Path from pathlib import Path
from setuptools import setup, find_packages from setuptools import setup, find_packages
from distutils.sysconfig import get_python_lib from distutils.sysconfig import get_python_lib
from torch.utils import cpp_extension, CUDA_HOME from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDA_HOME
os.environ["CC"] = "g++" os.environ["CC"] = "g++"
os.environ["CXX"] = "g++" os.environ["CXX"] = "g++"
...@@ -80,7 +80,7 @@ check_dependencies() ...@@ -80,7 +80,7 @@ check_dependencies()
arch_flags = get_compute_capabilities() arch_flags = get_compute_capabilities()
extensions = [ extensions = [
cpp_extension.CppExtension( CppExtension(
"awq_inference_engine", "awq_inference_engine",
[ [
"awq_cuda/pybind.cpp", "awq_cuda/pybind.cpp",
...@@ -96,7 +96,7 @@ extensions = [ ...@@ -96,7 +96,7 @@ extensions = [
additional_setup_kwargs = { additional_setup_kwargs = {
"ext_modules": extensions, "ext_modules": extensions,
"cmdclass": {'build_ext': cpp_extension.BuildExtension} "cmdclass": {'build_ext': BuildExtension}
} }
common_setup_kwargs.update(additional_setup_kwargs) common_setup_kwargs.update(additional_setup_kwargs)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment