def _linear_flops(module):
    """
    linear layer is a 2d matmul:
       out = (B, D) x (D, O)

    out is now of shape (B, O)
    """
    B, D = module.input_size
    D, O = module.weight.data.size()
    mult_flops = D * O * B
    add_flops = mult_flops
    comp_flops = 0
    return mult_flops, add_flops, comp_flops, 0, 0