class_docstr: ## Spatial Transformer
class_docstr_tok: ["#","#","Spatial","Transformer"]
docstr: :param channels: is the number of channels in the feature map
  :param n_heads: is the number of attention heads
  :param n_layers: is the number of transformer layers
  :param d_cond: is the size of the conditional embedding
docstr_tok: [":","param","channels",":","is","the","number","of","channels","in","the","feature","map",":","param","n_heads",":","is","the","number","of","attention","heads",":","param","n_layers",":","is","the","number","of","transformer","layers",":","param","d_cond",":","is","the","size","of","the","conditional","embedding"]
code: def __init__(
      self, channels: int, n_heads: int, n_layers: int, d_cond: int
  ):
      """
      :param channels: is the number of channels in the feature map
      :param n_heads: is the number of attention heads
      :param n_layers: is the number of transformer layers
      :param d_cond: is the size of the conditional embedding
      """
      super().__init__()
      # Initial group normalization
      self.norm = torch.nn.GroupNorm(
          num_groups=32, num_channels=channels, eps=1e-6, affine=True
      )
      # Initial $1 \times 1$ convolution
      self.proj_in = nn.Conv2d(
          channels, channels, kernel_size=1, stride=1, padding=0
      )
  
      # Transformer layers
      self.transformer_blocks = nn.ModuleList(
          [
              BasicTransformerBlock(
                  channels, n_heads, channels // n_heads, d_cond=d_cond
              )
              for _ in range(n_layers)
          ]
      )
  
      # Final $1 \times 1$ convolution
      self.proj_out = nn.Conv2d(
          channels, channels, kernel_size=1, stride=1, padding=0
      )
  
imports: from typing import Optional
  import torch
  import torch.nn.functional
  from torch import nn
label_desc: Use image node_id 1 to create a new SpatialTransformer object from inherited base classes: nn with example: obj = SpatialTransformer(channels, n_heads, n_layers, d_cond)