    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, representation_size=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0., hybrid_backbone=None, norm_layer=None,
                 num_frames=8, time_init='rand', attention_style='frozen_in_time'):
        """
        Args:
            img_size: (`int`, `tuple`)
                Input image size.
            patch_size: (`int`, `tuple`)
                Patch size.
            in_chans: (`int`)
                Number of input channels
            num_classes: (`int`)
                Number of classes for classification head.
            embed_dim: (`int`)
                Embedding dimension.
            depth: (`int`)
                Depth of transformer.
            num_heads: (`int`)
                Number of attention heads.
            mlp_ratio: (`int`)
                Ratio of mlp hidden dim to embedding dim.
            qkv_bias: (`bool`)
                Enable bias for qkv if True.
            qk_scale: (`float`)
                Override default qk scale of head_dim ** -0.5 if set.
            representation_size: (`Optional[int]`)
                Enable and set representation layer (pre-logits) to this value if set.
            drop_rate: (`float`)
                Dropout rate.
            attn_drop_rate: (`float`)
                Attention dropout rate.
            drop_path_rate: (`float`)
                Stochastic depth rate.
            hybrid_backbone: (`nn.Module`)
                CNN backbone to use in-place of PatchEmbed module.
            norm_layer: (`nn.Module`)
                Normalization layer.
            num_frames: (`int`)
                Maximum number of frames expected as input.
            time_init: (`str`)
                How to initialise the time attention layer, 'zeros' allows for the timesformer to start off
                as ViT.
            attention_style: (`str`)
                How to attend to space and time.
        """
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        self.attention_style = attention_style
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        logging.info('######USING ATTENTION STYLE: %s', self.attention_style)
        if hybrid_backbone is not None:
            raise NotImplementedError('hybrid backbone not implemented')
        else:
            # self.patch_embed = VideoPatchEmbed(
            #     img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, num_frames=num_frames)

            self.patch_embed = PatchEmbed2D(
                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, flatten=False)

        num_patches = self.patch_embed.num_patches*num_frames
        self.patches_per_frame = num_patches // num_frames

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.patches_per_frame + 1,
                        embed_dim))  # remember to take pos_embed[1:] for tiling over time
        if self.attention_style != 'bridge_former':
            self.temporal_embed = nn.Parameter(torch.zeros(1, num_frames, embed_dim))

        self.pos_drop = nn.Dropout(p=drop_rate)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            SpaceTimeBlock(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, time_init=time_init,
                attention_style=attention_style)
            for i in range(depth)])
        if self.attention_style == 'bridge_former':
            self.norm1 = norm_layer(embed_dim)
        else:
            self.norm = norm_layer(embed_dim)

        # Representation layer
        if representation_size:
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(OrderedDict([
                ('fc', nn.Linear(embed_dim, representation_size)),
                ('act', nn.Tanh())
            ]))
        else:
            self.pre_logits = nn.Identity()

        # Classifier head
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)

        # if num_frames > 1, then we perform ViT inflation and initialise time attention to zero so not necessary.
        if num_frames == 1:
            self.apply(init_vit_weights)

        # einops transformations
        self.einops_from_space = 'b (f n) d'
        self.einops_to_space = '(b f) n d'
        self.einops_from_time = 'b (f n) d'
        self.einops_to_time = '(b n) f d'
Args:
img_size: (`int`, `tuple`)
Input image size.
patch_size: (`int`, `tuple`)
Patch size.
in_chans: (`int`)
Number of input channels
num_classes: (`int`)
Number of classes for classification head.
embed_dim: (`int`)
Embedding dimension.
depth: (`int`)
Depth of transformer.
num_heads: (`int`)
Number of attention heads.
mlp_ratio: (`int`)
Ratio of mlp hidden dim to embedding dim.
qkv_bias: (`bool`)
Enable bias for qkv if True.
qk_scale: (`float`)
Override default qk scale of head_dim ** -0.5 if set.
representation_size: (`Optional[int]`)
Enable and set representation layer (pre-logits) to this value if set.
drop_rate: (`float`)
Dropout rate.
attn_drop_rate: (`float`)
Attention dropout rate.
drop_path_rate: (`float`)
Stochastic depth rate.
hybrid_backbone: (`nn.Module`)
CNN backbone to use in-place of PatchEmbed module.
norm_layer: (`nn.Module`)
Normalization layer.
num_frames: (`int`)
Maximum number of frames expected as input.
time_init: (`str`)
How to initialise the time attention layer, 'zeros' allows for the timesformer to start off
as ViT.
attention_style: (`str`)
How to attend to space and time.
towhee/models/frozen_in_time/frozen_video_transformer.py