class_docstr: Vision Transformer
  
  A PyTorch impl of : `Space-Time Transformer` from Frozen-in-time  - by Max Bain.
      https://arxiv.org/abs/2104.00650
  
  Based off:
   - ViT implementation from the timm library [https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py]
  lucidrains timesformer implementation [https://github.com/lucidrains/TimeSformer-pytorch].
  
  Notable differences:
   - allows for variable length input frames (<= num_frames)
   - allows for variable length input resolution  (<= (img_size, img_size)) [UNTESTED]
   - different attention block mechanism
class_docstr_tok: ["Vision","Transformer","A","PyTorch","impl","of",":","`","Space-Time","Transformer","`","from","Frozen-in-time","-","by","Max","Bain",".","https",":","\/\/arxiv.org\/abs\/2104.00650","Based","off",":","-","ViT","implementation","from","the","timm","library","[","https",":","\/\/github.com\/rwightman\/pytorch-image-models\/blob\/master\/timm\/models\/vision_transformer.py","]","lucidrains","timesformer","implementation","[","https",":","\/\/github.com\/lucidrains\/TimeSformer-pytorch","]",".","Notable","differences",":","-","allows","for","variable","length","input","frames","(","<","=","num_frames",")","-","allows","for","variable","length","input","resolution","(","<","=","(","img_size",",","img_size",")",")","[","UNTESTED","]","-","different","attention","block","mechanism"]
code: def forward_features(self, x):
      b, curr_frames, _, _, _ = x.shape
      x = self.patch_embed(
          x.transpose(2, 1)
      )  # [b*curr_frames, embed_dim, (original 244)14, (original 244)14]
      x = x.flatten(2).transpose(
          2, 1
      )  # [b*curr_frames,14*14, embed_dim]
      x = x.reshape(
          b, -1, self.embed_dim
      )  # [b,14*14*curr_frames,embed_dim]
      bf = x.shape[0]
      cls_tokens = self.cls_token.expand(
          bf, -1, -1
      )  # stole cls_tokens impl from Phil Wang, thanks
      x = torch.cat(
          (cls_tokens, x), dim=1
      )  # [b,14*14*curr_frames+1,embed_dim]
      # positional embed needs to be tiled for each frame (this does [1,2,3] --> [1,2,3,1,2,3]...)
      cls_embed = self.pos_embed[:, 0, :].unsqueeze(1)
      tile_pos_embed = self.pos_embed[:, 1:, :].repeat(
          1, self.num_frames, 1
      )
      if self.attention_style != "bridge_former":
          # temporal embed needs to be repeated within each frame (this does [1,2,3] --> [1,1,1,2,2,2,3,3,3]...)
          tile_temporal_embed = self.temporal_embed.repeat_interleave(
              self.patches_per_frame, 1
          )
          total_pos_embed = tile_pos_embed + tile_temporal_embed
      else:
          total_pos_embed = tile_pos_embed
      total_pos_embed = torch.cat([cls_embed, total_pos_embed], dim=1)
  
      curr_patches = x.shape[1]
      x = (
          x + total_pos_embed[:, :curr_patches]
      )  # [b,14*14*curr_frames+1,embed_dim]
      x = self.pos_drop(x)
      n = self.patches_per_frame
      f = curr_frames
  
      for blk in self.blocks:
          x = blk(
              x,
              self.einops_from_space,
              self.einops_to_space,
              self.einops_from_time,
              self.einops_to_time,
              time_n=n,
              space_f=f,
          )
      if self.attention_style == "bridge_former":
          x = self.norm1(x)[:, 0]
      else:
          x = self.norm(x)[:, 0]
      x = self.pre_logits(x)
  
      return x
  
imports: from collections import OrderedDict
  from functools import partial
  import torch
  from einops import rearrange, repeat
  from timm.models.layers import DropPath, trunc_normal_
  from torch import einsum, nn
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from towhee.models.utils.init_vit_weights import init_vit_weights
  import logging
label_desc: Use image node_id 5 for calling the SpaceTimeTransformer obj's underlying member method code with example usage: obj.forward_features(x) and returns: x