def flops_linear(module, input, output):
    """FLOPs estimation for `torch.nn.Linear`"""
    mm_flops = input.shape[0] * output.shape[1] * (2 * input.shape[1] - 1)
    bias_flops = output.numel() if module.bias is not None else 0
    return mm_flops + bias_flops
