rename module (#3)

mayank31398 · Aug 7, 2024 · b28f3e8 · b28f3e8
1 parent f545648
commit b28f3e8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/kernel_hyperdrive/utils.py b/kernel_hyperdrive/utils.py
@@ -5,7 +5,7 @@
 
 def compile_helpers() -> None:
     load_cpp_extension(
-        "vector_addition_cuda",
+        "khd",
         sources=[
             os.path.join(os.path.dirname(__file__), "vector_addition/cuda_kernel/vector_addition.cpp"),
             os.path.join(os.path.dirname(__file__), "vector_addition/cuda_kernel/vector_addition.cu"),

diff --git a/kernel_hyperdrive/vector_addition/cuda_kernel/__init__.py b/kernel_hyperdrive/vector_addition/cuda_kernel/__init__.py
@@ -6,9 +6,9 @@
 
 class _VectorAddition_CUDA(torch.autograd.Function):
     def forward(ctx, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        import vector_addition_cuda
+        import khd
 
-        return vector_addition_cuda.vector_addition_forward(x, y)
+        return khd.vector_addition_forward(x, y)
 
     def backward(ctx, output_grad: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         return output_grad, output_grad