limit chunk.axis

PaddlePaddle · Aug 6, 2021 · 977aa74 · 977aa74 · paddle-bot-old · Aug 6, 2021
1 parent 4cc3d9a
commit 977aa74
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
@@ -4890,6 +4890,7 @@ def split(input, num_or_sections, dim=-1, name=None):
         if isinstance(dim, Variable):
             dim = dim.numpy()
             dim = dim.item(0)
+        assert len(input.shape) + dim >= 0, "(rank(x) + axis) must >= 0"
         dim = (len(input.shape) + dim) if dim < 0 else dim
         attrs += ('axis', dim)
 
@@ -4951,6 +4952,7 @@ def _get_SectionsTensorList(one_list):
         dim.stop_gradient = True
         inputs['AxisTensor'] = dim
     else:
+        assert len(input.shape) + dim >= 0, "(rank(x) + axis) must >= 0"
         dim = (len(input_shape) + dim) if dim < 0 else dim
         attrs['axis'] = dim