text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from .blocks import *
  
  
  class UNet2D(nn.Module):
      def __init__(self, in_channels, out_channels, conv_depths=(64, 128, 256, 512, 1024)):
          assert len(conv_depths) > 2, 'conv_depths must have at least 3 members'
  
          super(UNet2D, self).__init__()
  
          # defining encoder layers
          encoder_layers = []
          encoder_layers.append(First2D(in_channels, conv_depths[0], conv_depths[0]))
          encoder_layers.extend([Encoder2D(conv_depths[i], conv_depths[i + 1], conv_depths[i + 1])
                                 for i in range(len(conv_depths)-2)])
  
          # defining decoder layers
          decoder_layers = []
          decoder_layers.extend([Decoder2D(2 * conv_depths[i + 1], 2 * conv_depths[i], 2 * conv_depths[i], conv_depths[i])
                                 for i in reversed(range(len(conv_depths)-2))])
          decoder_layers.append(Last2D(conv_depths[1], conv_depths[0], out_channels))
  
          # encoder, center and decoder layers
          self.encoder_layers = nn.Sequential(*encoder_layers)
          self.center = Center2D(conv_depths[-2], conv_depths[-1], conv_depths[-1], conv_depths[-2])
          self.decoder_layers = nn.Sequential(*decoder_layers)
  
      def forward(self, x, return_all=False):
          x_enc = [x]
          for enc_layer in self.encoder_layers:
              x_enc.append(enc_layer(x_enc[-1]))
  
          x_dec = [self.center(x_enc[-1])]
          for dec_layer_idx, dec_layer in enumerate(self.decoder_layers):
              x_opposite = x_enc[-1-dec_layer_idx]
              x_cat = torch.cat(
                  [pad_to_shape(x_dec[-1], x_opposite.shape), x_opposite],
                  dim=1
              )
              x_dec.append(dec_layer(x_cat))
  
          if not return_all:
              return x_dec[-1]
          else:
              return x_enc + x_dec
  
  
  class UNet3D(nn.Module):
      def __init__(self, in_channels, out_channels, conv_depths=(64, 128, 256, 512, 1024)):
          assert len(conv_depths) > 2, 'conv_depths must have at least 3 members'
  
          super(UNet3D, self).__init__()
  
          # defining encoder layers
          encoder_layers = []
          encoder_layers.append(First3D(in_channels, conv_depths[0], conv_depths[0]))
          encoder_layers.extend([Encoder3D(conv_depths[i], conv_depths[i + 1], conv_depths[i + 1])
                                 for i in range(len(conv_depths)-2)])
  
          # defining decoder layers
          decoder_layers = []
          decoder_layers.extend([Decoder3D(2 * conv_depths[i + 1], 2 * conv_depths[i], 2 * conv_depths[i], conv_depths[i])
                                 for i in reversed(range(len(conv_depths)-2))])
          decoder_layers.append(Last3D(conv_depths[1], conv_depths[0], out_channels))
  
          # encoder, center and decoder layers
          self.encoder_layers = nn.Sequential(*encoder_layers)
          self.center = Center3D(conv_depths[-2], conv_depths[-1], conv_depths[-1], conv_depths[-2])
          self.decoder_layers = nn.Sequential(*decoder_layers)
  
      def forward(self, x, return_all=False):
          x_enc = [x]
          for enc_layer in self.encoder_layers:
              x_enc.append(enc_layer(x_enc[-1]))
  
          x_dec = [self.center(x_enc[-1])]
          for dec_layer_idx, dec_layer in enumerate(self.decoder_layers):
              x_opposite = x_enc[-1-dec_layer_idx]
              x_cat = torch.cat(
                  [pad_to_shape(x_dec[-1], x_opposite.shape), x_opposite],
                  dim=1
              )
              x_dec.append(dec_layer(x_cat))
  
          if not return_all:
              return x_dec[-1]
          else:
              return x_enc + x_dec
  
  
  def pad_to_shape(this, shp):
      """
      Pads this image with zeroes to shp.
      Args:
          this: image tensor to pad
          shp: desired output shape
  
      Returns:
          Zero-padded tensor of shape shp.
      """
      if len(shp) == 4:
          pad = (0, shp[3] - this.shape[3], 0, shp[2] - this.shape[2])
      elif len(shp) == 5:
          pad = (0, shp[4] - this.shape[4], 0, shp[3] - this.shape[3], 0, shp[2] - this.shape[2])
      return F.pad(this, pad)
  