code: def __init__(
      self,
      *,
      vit,
      num_memories_per_layer=10,
      num_classes=2,
  ):
      super().__init__()
      assert isinstance(vit, ViT)
  
      # extract some model variables needed
  
      dim = vit.cls_token.shape[-1]
      layers = len(vit.transformer.layers)
      num_patches = vit.pos_embedding.shape[-2]
  
      self.vit = vit
  
      # freeze ViT backbone - only memories will be finetuned
  
      freeze_all_layers_(vit)
  
      # learnable parameters
  
      self.memory_cls_token = nn.Parameter(torch.randn(dim))
      self.memories_per_layer = nn.Parameter(
          torch.randn(layers, num_memories_per_layer, dim)
      )
  
      self.mlp_head = nn.Sequential(
          nn.LayerNorm(dim), nn.Linear(dim, num_classes)
      )
  
      # specialized attention mask to preserve the output of the original ViT
      # it allows the memory CLS token to attend to all other tokens (and the learnable memory layer tokens), but not vice versa
  
      attn_mask = torch.ones(
          (num_patches, num_patches), dtype=torch.bool
      )
      attn_mask = F.pad(
          attn_mask, (1, num_memories_per_layer), value=False
      )  # main tokens cannot attend to learnable memories per layer
      attn_mask = F.pad(
          attn_mask, (0, 0, 1, 0), value=True
      )  # memory CLS token can attend to everything
      self.register_buffer("attn_mask", attn_mask)
  
imports: import torch
  from torch import nn
  import torch.nn.functional
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange
label_desc: Use image node_id 1 to create a new Adapter object from inherited base classes: nn with example: obj = Adapter()