class_docstr: ## Spatial Transformer
class_docstr_tok: ["#","#","Spatial","Transformer"]
docstr: :param x: is the feature map of shape `[batch_size, channels, height, width]`
  :param cond: is the conditional embeddings of shape `[batch_size,  n_cond, d_cond]`
docstr_tok: [":","param","x",":","is","the","feature","map","of","shape","`","[","batch_size",",","channels",",","height",",","width","]","`",":","param","cond",":","is","the","conditional","embeddings","of","shape","`","[","batch_size",",","n_cond",",","d_cond","]","`"]
code: def forward(self, x: torch.Tensor, cond: torch.Tensor):
      """
      :param x: is the feature map of shape `[batch_size, channels, height, width]`
      :param cond: is the conditional embeddings of shape `[batch_size,  n_cond, d_cond]`
      """
      # Get shape `[batch_size, channels, height, width]`
      b, c, h, w = x.shape
      # For residual connection
      x_in = x
      # Normalize
      x = self.norm(x)
      # Initial $1 \times 1$ convolution
      x = self.proj_in(x)
      # Transpose and reshape from `[batch_size, channels, height, width]`
      # to `[batch_size, height * width, channels]`
      x = x.permute(0, 2, 3, 1).view(b, h * w, c)
      # Apply the transformer layers
      for block in self.transformer_blocks:
          x = block(x, cond)
      # Reshape and transpose from `[batch_size, height * width, channels]`
      # to `[batch_size, channels, height, width]`
      x = x.view(b, h, w, c).permute(0, 3, 1, 2)
      # Final $1 \times 1$ convolution
      x = self.proj_out(x)
      # Add residual
      return x + x_in
  
imports: from typing import Optional
  import torch
  import torch.nn.functional
  from torch import nn
label_desc: Use image node_id 2 for calling the SpatialTransformer obj's underlying member method code with example usage: obj.forward(x, cond) and returns: unknown