def num_params(model, count_fixed=False, display_all_modules=False,
               print_stats=False):
    """Counts number of parameters and layers in the model

    _exclude_from_layer_count contains names that shouldn't be counted
    as layers, shortcut layer, for example.

    Args:
        model (nn.Module):
        count_fixed (bool): if True, modules with fixed params
            (.requires_grad=False) are counted toward total_num_params too
        display_all_modules (bool): if True, prints info on all of the modules

    Returns:
        total_num_params, total_num_layers
    """
    _exclude_from_layer_count = ['shortcut']
    total_num_params = 0
    total_num_layers = 0
    for n, p in model.named_parameters():
        if p.requires_grad or count_fixed:
            num_params = 1
            for s in p.shape:
                num_params *= s
            if len(p.data.size()) > 1:
                excl_in_name = len(list(filter(lambda excl: excl in n, _exclude_from_layer_count))) > 0
                if not excl_in_name:
                    total_num_layers += 1
        else:
            num_params = 0
        if print_stats and display_all_modules: print("{}: {}".format(n, num_params))
        total_num_params += num_params
    if print_stats:
        print("-" * 50)
        print("Total number of parameters: {:.2e}".format(total_num_params))
        print("-" * 50)
        print("Total number of layers: {}".format(total_num_layers))
    return total_num_params, total_num_layers