class_docstr: Class Attention Layer as in CaiT https://arxiv.org/abs/2103.17239
class_docstr_tok: ["Class","Attention","Layer","as","in","CaiT","https",":","\/\/arxiv.org\/abs\/2103.17239"]
code: def __init__(
      self,
      dim,
      num_heads,
      mlp_ratio=4.0,
      qkv_bias=False,
      qk_scale=None,
      drop=0.0,
      attn_drop=0.0,
      drop_path=0.0,
      act_layer=nn.GELU,
      norm_layer=nn.LayerNorm,
      eta=None,
      tokens_norm=False,
  ):
      super().__init__()
      self.norm1 = norm_layer(dim)
  
      self.attn = ClassAttention(
          dim,
          num_heads=num_heads,
          qkv_bias=qkv_bias,
          qk_scale=qk_scale,
          attn_drop=attn_drop,
          proj_drop=drop,
      )
  
      self.drop_path = (
          DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
      )
      self.norm2 = norm_layer(dim)
      mlp_hidden_dim = int(dim * mlp_ratio)
      self.mlp = Mlp(
          in_features=dim,
          hidden_features=mlp_hidden_dim,
          act_layer=act_layer,
          drop=drop,
      )
  
      if (
          eta is not None
      ):  # LayerScale Initialization (no layerscale when None)
          self.gamma1 = nn.Parameter(
              eta * torch.ones(dim), requires_grad=True
          )
          self.gamma2 = nn.Parameter(
              eta * torch.ones(dim), requires_grad=True
          )
      else:
          self.gamma1, self.gamma2 = 1.0, 1.0
  
      # FIXME: A hack for models pre-trained with layernorm over all the tokens not just
      # the CLS
      self.tokens_norm = tokens_norm
  
imports: import logging
  import math
  from functools import partial
  from typing import List, Optional
  import torch
  import torch.nn
  from fairscale.nn import checkpoint_wrapper
  from vissl.config import AttrDict
  from vissl.models.model_helpers import DropPath, to_2tuple, trunc_normal_
  from vissl.models.trunks import register_model_trunk
  from vissl.models.trunks.vision_transformer import Mlp
  from vissl.utils.fsdp_utils import fsdp_wrapper
  from vissl.utils.misc import set_torch_seed
label_desc: Use image node_id 1 to create a new ClassAttentionBlock object from inherited base classes: nn with example: obj = ClassAttentionBlock(dim, num_heads, mlp_ratio, qkv_bias, qk_scale, drop, attn_drop, drop_path, act_layer, norm_layer, eta, tokens_norm)