class_docstr: Attention module that can take tensor with [B, N, C] or [B, C, H, W] as input. Structure specified by metaformer
  Modified from:
  https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
  Args:
      dim (int): dimension
      head_dim (int): head dimension
class_docstr_tok: ["Attention","module","that","can","take","tensor","with","[","B",",","N",",","C","]","or","[","B",",","C",",","H",",","W","]","as","input",".","Structure","specified","by","metaformer","Modified","from",":","https",":","\/\/github.com\/rwightman\/pytorch-image-models\/blob\/master\/timm\/models\/vision_transformer.py","Args",":","dim","(","int",")",":","dimension","head_dim","(","int",")",":","head","dimension"]
code: def __init__(
      self,
      dim,
      head_dim=32,
      qkv_bias=False,
      attn_drop=0.0,
      proj_drop=0.0,
  ):
      super().__init__()
      assert dim % head_dim == 0, "dim should be divisible by head_dim"
      self.head_dim = head_dim
      self.num_heads = dim // head_dim
      self.scale = head_dim**-0.5
  
      self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj = nn.Linear(dim, dim)
      self.proj_drop = nn.Dropout(proj_drop)
  
imports: from torch import nn
  import torch
label_desc: Use image node_id 1 to create a new Attention object from inherited base classes: nn with example: obj = Attention(dim, head_dim, qkv_bias, attn_drop, proj_drop)