code: def __init__(
      self,
      dim,
      fmap_size,
      depth,
      heads,
      dim_key,
      dim_value,
      mlp_mult=2,
      dropout=0.0,
      dim_out=None,
      downsample=False,
  ):
      super().__init__()
      dim_out = default(dim_out, dim)
      self.layers = nn.ModuleList([])
      self.attn_residual = (not downsample) and dim == dim_out
  
      for _ in range(depth):
          self.layers.append(
              nn.ModuleList(
                  [
                      Attention(
                          dim,
                          fmap_size=fmap_size,
                          heads=heads,
                          dim_key=dim_key,
                          dim_value=dim_value,
                          dropout=dropout,
                          downsample=downsample,
                          dim_out=dim_out,
                      ),
                      FeedForward(dim_out, mlp_mult, dropout=dropout),
                  ]
              )
          )
  
imports: from math import ceil
  import torch
  from torch import nn, einsum
  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 Transformer object from inherited base classes: nn with example: obj = Transformer(dim, fmap_size, depth, heads, dim_key, dim_value, mlp_mult, dropout, dim_out, downsample)