From 62041b78bfaad78267556a297d3f25b6cebf3152 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 05:20:44 +0000 Subject: [PATCH] fix: import torch.utils.cpp_extension explicitly in production_compress --- dsv4/kernels/compressor/production_compress.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dsv4/kernels/compressor/production_compress.py b/dsv4/kernels/compressor/production_compress.py index 8fd37126..1060f386 100644 --- 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]"],