import torch


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