import torch


def unrollout(tensor, n_channels):
    """
    tensor: B * C x ... -> B x C x ...
    """
    shape = tensor.shape
    new_shape = [shape[0] // n_channels, n_channels]
    if len(shape) > 1:
        new_shape += shape[1:]
    return torch.reshape(tensor, new_shape)