def flatten(tensor):
    """
    Flatten input tensor as the shape of (nb, nf)

    :param tensor: input Tensor
    :type tensor: torch.Tensor
    :return:  flattened tensor
    :rtype: torch.Tensor
    """
    assert len(tensor.shape) >= 2

    if len(tensor.shape) > 2:
        flattened = tensor.view(tensor.shape[0], -1)
    else:
        flattened = tensor

    return flattened