content: import torch
  from fvcore.nn import FlopCountAnalysis, flop_count_table
  from torch.nn import MultiheadAttention
  
  from models.beit.st_beit import BeitConfig, BeitModel
  from models.temporal_model import (STAdapter, TemporalAttention,
                                     WindowTemporalAttention)
  
  
  def mem_stat():
      mem = torch.cuda.max_memory_allocated() / 1024 / 1024
      print(f"max memory allocated: {mem}MB")
  
  
  def build_backbone(tm_block="timesformer"):
      """TODO: Docstring for build_backbone.
      Returns: TODO
  
      """
      if tm_block == "timesformer":
          other_cfg = dict(
              num_frames=12, temporal_model_block="timesformer", temporal_model_config={}
          )
      elif tm_block == "st_adapter":
          other_cfg = dict(
              num_frames=12, temporal_model_block="st_adapter", temporal_model_config={}
          )
      elif tm_block == "xclip":
          other_cfg = dict(
              num_frames=12, temporal_model_block="xclip", temporal_model_config={}
          )
      elif tm_block == "none":
          other_cfg = dict(num_frames=12, temporal_model_block="none", temporal_model_config={})
      elif tm_block == "wa_2x2":
          other_cfg = dict(
              num_frames=12,
              temporal_model_block="window_attention",
              temporal_model_config=dict(window_size=(2, 2)),
          )
      elif tm_block == "wa_7x7":
          other_cfg = dict(
              num_frames=12,
              temporal_model_block="window_attention",
              temporal_model_config=dict(window_size=(7, 7)),
          )
      else:
          raise ValueError("not exist")
  
      model_card = "microsoft/beit-base-patch16-224-pt22k-ft22k"
      model_config = BeitConfig.from_pretrained(model_card, image_size=224, **other_cfg)
      model = BeitModel(model_config)
      return model
  
  
  # model = TemporalAttention()
  model = build_backbone("st_adapter")
  model.gradient_checkpointing_enable()
  model.cuda()
  for i in range(3):
      x = torch.rand(32, 12, 3, 224, 224, requires_grad=True)
      x = x.cuda()
      x = x.requires_grad_()
      y = model(x)
      loss = y[0].mean()
      loss.backward()
      mem_stat()
  
  # flops = FlopCountAnalysis(model, x)
  # print(flop_count_table(flops))
  