def get_visual_dim_in(
    trans_func,
    width_per_group,
    pooling,
):
    """
    Compute the input dimension to the VisualClassifyHead.
    args:
        trans_func (string): transform function to be used to contrusct each
            ResBlock.
        width_per_group (int): width of each group.
        pooling (bool): Whether to use the output of the ConvNet.
    returns:
        dim_in (int or list): the input dimension.
    """
    if trans_func == "basic_transform":
        factor = 2 ** 3
    elif trans_func == "bottleneck_transform":
        factor = 4 * (2 ** 3)
    else:
        raise NotImplementedError(
            "Does not support {} transfomration".format(trans_func)
        )
    dim_in = [width_per_group * factor]

    if pooling:
        dim_in = sum(dim_in)

    return dim_in