Yeah, I think a single function would be nicer. Both tf.split and numpy.split have that sort of API.

Thanks for the feedback, I will merge it in torch.split then.

I merged both functions now.

Before that, I timed all functions on my machine (tensor=torch.randn(200, 10, 2, 2) into 40 chunks of size 5 in dim=0)

• merged function with split_size_or_sections = [5] * 40: ~8.1586s
• merged function with split_size_or_sections = 5: ~7.4982
• native with `split_size=5: ~7.2688s

What are your thoughts? Any suggestion on the code / naming / documentation?

@pytorchbot test this please

@ptrblck great initiative, I've been using my own wrapper for a while. Would be nice to have it in the code base.

You could simplify the code quite a bit though by using plain python operations instead of invoking torch overheads, e.g. something like

def split(tensor, sizes, dim=0):
    if dim < 0:
        dim += tensor.dim()

    if isinstance(sizes, int):
        # original code ...
        return chunks

    if tensor.size(dim) != sum(sizes):
        raise ValueError("Sizes do not match tensor size in dim")

    nsizes = len(sizes)
    sizes = [0] + sizes
    return tuple(tensor.narrow(dim, sizes[i], sizes[i + 1])
                 for i in range(nsizes))

Should be slightly faster too.

@ptrblck as soon as you add unit tests for the list of splits case, i can merge this in.

@flennerhag Thanks for the suggestions! I tried to change some Pytorch code to plain python operations.
@soumith I also added some tests in test_split. If it's not sufficient, I can add some more test cases.

@pytorchbot test this please

thanks a lot @ptrblck !