fix: import torch.utils.cpp_extension explicitly in production_compress

2026-06-01 05:20:44 +00:00
parent 2155fd6c90
commit 62041b78bf
1 changed files with 2 additions and 1 deletions
--- a/dsv4/kernels/compressor/production_compress.py
+++ b/dsv4/kernels/compressor/production_compress.py
@@ -22,8 +22,9 @@ def _get_kernel():
    global _kernel_module
    if _kernel_module is not None:
        return _kernel_module
+    from torch.utils.cpp_extension import load
    kernel_dir = os.path.join(os.path.dirname(__file__), "..", "cuda")
-    _kernel_module = torch.utils.cpp_extension.load(
+    _kernel_module = load(
        name="compressor_reduce",
        sources=[os.path.join(kernel_dir, "compressor_reduce.cu")],
        extra_cuda_cflags=["-O3", "--generate-code=arch=compute_100a,code=[sm_100a]"],