Solved by aRI0U in post #3 
              
              
                Thanks a lot for your very precise answer ! 
I modified it a bit in order not to use Numpy and it works great: 
def get_split_wo_numpy(x):
    d = x[:,-1].to(torch.int32)
    _, counts = torch.unique(d, return_counts=True)
    return torch.split(x, counts.tolist())

# create tensor
N = 50
c0 = torch…