diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py
index 8150cd6d9a7aa..08b99e5bf0c41 100644
--- a/python/tvm/relay/op/nn/nn.py
+++ b/python/tvm/relay/op/nn/nn.py
@@ -3157,7 +3157,8 @@ def correlation(
 
 
 def space_to_batch_nd(data, block_shape, paddings):
-    r"""Divide spatial dimensions of the data into a grid of blocks and interleave them into batch dim.
+    r"""Divide spatial dimensions of the data into a grid of blocks
+    and interleave them into batch dim.
 
     Parameters
     ----------
diff --git a/python/tvm/topi/testing/batch_to_space_nd.py b/python/tvm/topi/testing/batch_to_space_nd.py
index 108da6d970d0b..80af79b8cacb6 100644
--- a/python/tvm/topi/testing/batch_to_space_nd.py
+++ b/python/tvm/topi/testing/batch_to_space_nd.py
@@ -85,7 +85,7 @@ def batch_to_space_nd_python(data, block_shape, crop_begin_list, crop_end_list):
 
     for i, _ in enumerate(r_p_shape):
         strides.append(1)
-        if i > 0 and i <= M:
+        if 0 < i <= M:
             # begin and end index for spatial dimensions
             begin_idx.append(crop_begin_list[i - 1])
             end_idx.append(r_p_shape[i] - crop_end_list[i - 1])