class_docstr: ## U-Net model
docstr: :param x: is the input feature map of shape `[batch_size, channels, width, height]`
  :param time_steps: are the time steps of shape `[batch_size]`
  :param cond: conditioning of shape `[batch_size, n_cond, d_cond]`
docstr_tok: [":","param","x",":","is","the","input","feature","map","of","shape","`","[","batch_size",",","channels",",","width",",","height","]","`",":","param","time_steps",":","are","the","time","steps","of","shape","`","[","batch_size","]","`",":","param","cond",":","conditioning","of","shape","`","[","batch_size",",","n_cond",",","d_cond","]","`"]
code: def forward(
      self,
      x: torch.Tensor,
      time_steps: torch.Tensor,
      cond: torch.Tensor,
  ):
      """
      :param x: is the input feature map of shape `[batch_size, channels, width, height]`
      :param time_steps: are the time steps of shape `[batch_size]`
      :param cond: conditioning of shape `[batch_size, n_cond, d_cond]`
      """
      # To store the input half outputs for skip connections
      x_input_block = []
  
      # Get time step embeddings
      t_emb = self.time_step_embedding(time_steps)
      t_emb = self.time_embed(t_emb)
  
      # Input half of the U-Net
      for module in self.input_blocks:
          x = module(x, t_emb, cond)
          x_input_block.append(x)
      # Middle of the U-Net
      x = self.middle_block(x, t_emb, cond)
      # Output half of the U-Net
      for module in self.output_blocks:
          x = torch.cat([x, x_input_block.pop()], dim=1)
          x = module(x, t_emb, cond)
  
      # Final normalization and $3 \times 3$ convolution
      return self.out(x)
  
imports: import math
  from typing import List
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from labml_nn.diffusion.stable_diffusion.model.unet_attention import SpatialTransformer
label_desc: Use image node_id 3 for calling the UNetModel obj's underlying member method code with example usage: obj.forward(x, time_steps, cond) and returns: self