src/transformers/models/flava/modeling_flava.py/FlavaImageCodebookLayerGroup/__init__
class FlavaImageCodebookLayerGroup:    def __init__(self, num_blocks: int, num_layers: int, in_size: int, out_size: int, use_pool: bool = True):
        super().__init__()
        blocks = OrderedDict()
        for i in range(num_blocks):
            if i == 0:
                blocks[f"block_{i+1}"] = FlavaImageCodebookBlock(in_size, out_size, num_layers)
            else:
                blocks[f"block_{i+1}"] = FlavaImageCodebookBlock(out_size, out_size, num_layers)

        if use_pool:
            blocks["pool"] = nn.MaxPool2d(kernel_size=2)

        self.group = nn.Sequential(blocks)