From fb7c687548b01ca53ff93d1ce1177b46b7303470 Mon Sep 17 00:00:00 2001
From: dan_the_3rd <43445237+danthe3rd@users.noreply.github.com>
Date: Tue, 29 Jul 2025 10:14:15 +0200
Subject: [PATCH] Merge pull request #135 from danthe3rd/patch-3

Fix import-time error: "undefined symbol: _ZN3c104cuda20getCurrentCUDAStreamEa"
---
 setup.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 916f760..22a65ba 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@ import shutil
 import subprocess
 from setuptools import find_packages
 from setuptools.command.build_py import build_py
-from torch.utils.cpp_extension import CppExtension, CUDA_HOME
+from torch.utils.cpp_extension import CUDAExtension, CUDA_HOME
 
 current_dir = os.path.dirname(os.path.realpath(__file__))
 cxx_flags = ['-std=c++20', '-O3', '-fPIC', '-Wno-psabi']
@@ -89,7 +89,7 @@ if __name__ == '__main__':
             ]
         },
         ext_modules=[
-            CppExtension(name='deep_gemm_cpp',
+            CUDAExtension(name='deep_gemm_cpp',
                          sources=sources,
                          include_dirs=build_include_dirs,
                          libraries=build_libraries,