src/transformers/models/cvt/modeling_cvt.py/CvtStage/__init__
class CvtStage:    def __init__(self, config, stage):
        super().__init__()
        self.config = config
        self.stage = stage
        if self.config.cls_token[self.stage]:
            self.cls_token = nn.Parameter(torch.randn(1, 1, self.config.embed_dim[-1]))

        self.embedding = CvtEmbeddings(
            patch_size=config.patch_sizes[self.stage],
            stride=config.patch_stride[self.stage],
            num_channels=config.num_channels if self.stage == 0 else config.embed_dim[self.stage - 1],
            embed_dim=config.embed_dim[self.stage],
            padding=config.patch_padding[self.stage],
            dropout_rate=config.drop_rate[self.stage],
        )

        drop_path_rates = [x.item() for x in torch.linspace(0, config.drop_path_rate[self.stage], config.depth[stage])]

        self.layers = nn.Sequential(
            *[
                CvtLayer(
                    num_heads=config.num_heads[self.stage],
                    embed_dim=config.embed_dim[self.stage],
                    kernel_size=config.kernel_qkv[self.stage],
                    padding_q=config.padding_q[self.stage],
                    padding_kv=config.padding_kv[self.stage],
                    stride_kv=config.stride_kv[self.stage],
                    stride_q=config.stride_q[self.stage],
                    qkv_projection_method=config.qkv_projection_method[self.stage],
                    qkv_bias=config.qkv_bias[self.stage],
                    attention_drop_rate=config.attention_drop_rate[self.stage],
                    drop_rate=config.drop_rate[self.stage],
                    drop_path_rate=drop_path_rates[self.stage],
                    mlp_ratio=config.mlp_ratio[self.stage],
                    with_cls_token=config.cls_token[self.stage],
                )
                for _ in range(config.depth[self.stage])
            ]
        )