import torch


def get_output_dim(model, pooling_type="gem"):
    """Dinamically compute the output size of a model.
    """
    output_dim = model(torch.ones([2, 3, 224, 224])).shape[1]
    if pooling_type == "netvlad":
        output_dim *= 64  # NetVLAD layer has 64x bigger output dimensions
    return output_dim