diff --git a/spconv/pytorch/pool.py b/spconv/pytorch/pool.py
index 2384507..5659993 100644
--- a/spconv/pytorch/pool.py
+++ b/spconv/pytorch/pool.py
@@ -314,8 +314,10 @@ def __init__(self,
         self.record_voxel_count = record_voxel_count
         self.dilation = expand_nd(ndim, dilation)
         self.indice_key = indice_key
-        kv = int(np.prod(kernel_size))
-        assert kv <= 32, "avg pool only support implicit-gemm style indice gen with kv <= 32 limit"
+        kv = int(np.prod(self.kernel_size))
+        assert kv <= 128, "avg pool only support implicit-gemm style indice gen with kv <= 128 limit"
+        if kv >= 32:
+            pass # maybe show some warnings here
         self.algo = ConvAlgo.MaskImplicitGemm
 
     def extra_repr(self):
@@ -581,4 +583,4 @@ def __init__(self,
 
 ALL_POOL_LAYERS = set([
     SparseAvgPool3d, SparseAvgPool2d, SparseAvgPool1d, SparseMaxPool1d, SparseMaxPool2d, SparseMaxPool3d, SparseMaxPool4d, SparseAvgPool, SparseMaxPool
-])
\ No newline at end of file
+])