class_docstr: ## U-Net
class_docstr_tok: ["#","#","U-Net"]
docstr: * `x` has shape `[batch_size, in_channels, height, width]`
  * `t` has shape `[batch_size]`
code: def forward(self, x: torch.Tensor, t: torch.Tensor):
      """
      * `x` has shape `[batch_size, in_channels, height, width]`
      * `t` has shape `[batch_size]`
      """
  
      # Get time-step embeddings
      t = self.time_emb(t)
  
      # Get image projection
      x = self.image_proj(x)
  
      # `h` will store outputs at each resolution for skip connection
      h = [x]
      # First half of U-Net
      for m in self.down:
          x = m(x, t)
          h.append(x)
  
      # Middle (bottom)
      x = self.middle(x, t)
  
      # Second half of U-Net
      for m in self.up:
          if isinstance(m, Upsample):
              x = m(x, t)
          else:
              # Get the skip connection from first half of U-Net and concatenate
              s = h.pop()
              x = torch.cat((x, s), dim=1)
              #
              x = m(x, t)
  
      # Final normalization and convolution
      return self.final(self.act(self.norm(x)))
  
imports: import math
  from typing import Optional, Tuple, Union, List
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 2 for calling the UNet obj's underlying member method code with example usage: obj.forward(x, t) and returns: self