Hi,
I frequently encounter the situation where I have to split a tensor in non regular sub-parts then apply different operations on each part, and concatenate all the results together.
I found two approaches to do that and I wonder which one is better.
The following snippet of code attempts to present to the approaches.
def with_cat(input, dimensions, functions):
    outputs = []
    i = 0
    for dim, fn in zip(dimensions, functions):
        x = fn(input[:, i:i + dim])  # split and process
        outputs.append(x)
        i += dim

    return torch.cat(outputs, dim=1)  # merge


def without_cat(input, dimensions, functions):
    output = torch.empty_like(input)
    i = 0
    for dim, fn in zip(dimensions, functions):
        x = fn(input[:, i:i + dim])  # split and process
        output[:, i:i + dim] = x  # merge
        i += dim

    return output

note: both works similarly for forward and backward.
Which one is better ? In term of memory, execution time, idiomatic, …
Is there another better way to split process merge data ?
Thanks