class_docstr: The half UNet model with attention and timestep embedding.
  For usage, see UNet.
  
  Arguments
  ---------
  in_channels: int
      channels in the input Tensor.
  model_channels: int
      base channel count for the model.
  out_channels: int
      channels in the output Tensor.
  num_res_blocks: int
      number of residual blocks per downsample.
  attention_resolutions: int
      a collection of downsample rates at which
      attention will take place. May be a set, list, or tuple.
      For example, if this contains 4, then at 4x downsampling, attention
      will be used.
  dropout: float
      the dropout probability.
  channel_mult: int
      channel multiplier for each level of the UNet.
  conv_resample: bool
      if True, use learned convolutions for upsampling and
      downsampling
  emb_dim: int
      time embedding dimension (defaults to model_channels * 4)
  cond_emb: dict
      embeddings on which the model will be conditioned
  
      Example:
      {
          "speaker": {
              "emb_dim": 256
          },
          "label": {
              "emb_dim": 12
          }
      }
  use_cond_emb: dict
      a dictionary with keys corresponding to keys in cond_emb
      and values corresponding to Booleans that turn embeddings
      on and off. This is useful in combination with hparams files
      to turn embeddings on and off with simple switches
  
      Example:
      {"speaker": False, "label": True}
  dims: int
      determines if the signal is 1D, 2D, or 3D.
  num_heads: int
      the number of attention heads in each attention layer.
  num_heads_channels: int
      if specified, ignore num_heads and instead use
                             a fixed channel width per attention head.
  num_heads_upsample: int
      works with num_heads to set a different number
                             of heads for upsampling. Deprecated.
  resblock_updown: bool
      use residual blocks for up/downsampling.
  
  use_fixup_init: bool
      whether to use FixUp initialization
  
  Example
  -------
  >>> model = DecoderUNetModel(
  ...    in_channels=1,
  ...    model_channels=32,
  ...    out_channels=3,
  ...    num_res_blocks=1,
  ...    attention_resolutions=[1]
  ... )
  >>> x = torch.randn(4, 1, 2, 4)
  >>> ts = torch.tensor([10, 100, 50, 25])
  >>> out = model(x, ts)
  >>> out.shape
  torch.Size([4, 3, 16, 32])
class_docstr_tok: ["The","half","UNet","model","with","attention","and","timestep","embedding",".","For","usage",",","see","UNet",".","Arguments","--","--","--","--","-","in_channels",":","int","channels","in","the","input","Tensor",".","model_channels",":","int","base","channel","count","for","the","model",".","out_channels",":","int","channels","in","the","output","Tensor",".","num_res_blocks",":","int","number","of","residual","blocks","per","downsample",".","attention_resolutions",":","int","a","collection","of","downsample","rates","at","which","attention","will","take","place",".","May","be","a","set",",","list",",","or","tuple",".","For","example",",","if","this","contains","4",",","then","at","4x","downsampling",",","attention","will","be","used",".","dropout",":","float","the","dropout","probability",".","channel_mult",":","int","channel","multiplier","for","each","level","of","the","UNet",".","conv_resample",":","bool","if","True",",","use","learned","convolutions","for","upsampling","and","downsampling","emb_dim",":","int","time","embedding","dimension","(","defaults","to","model_channels","*","4",")","cond_emb",":","dict","embeddings","on","which","the","model","will","be","conditioned","Example",":","{","``","speaker","''",":","{","``","emb_dim","''",":","256","}",",","``","label","''",":","{","``","emb_dim","''",":","12","}","}","use_cond_emb",":","dict","a","dictionary","with","keys","corresponding","to","keys","in","cond_emb","and","values","corresponding","to","Booleans","that","turn","embeddings","on","and","off",".","This","is","useful","in","combination","with","hparams","files","to","turn","embeddings","on","and","off","with","simple","switches","Example",":","{","``","speaker","''",":","False",",","``","label","''",":","True","}","dims",":","int","determines","if","the","signal","is","1D",",","2D",",","or","3D",".","num_heads",":","int","the","number","of","attention","heads","in","each","attention","layer",".","num_heads_channels",":","int","if","specified",",","ignore","num_heads","and","instead","use","a","fixed","channel","width","per","attention","head",".","num_heads_upsample",":","int","works","with","num_heads","to","set","a","different","number","of","heads","for","upsampling",".","Deprecated",".","resblock_updown",":","bool","use","residual","blocks","for","up\/downsampling",".","use_fixup_init",":","bool","whether","to","use","FixUp","initialization","Example","--","--","--","-",">",">",">","model","=","DecoderUNetModel","(","...","in_channels=1",",","...","model_channels=32",",","...","out_channels=3",",","...","num_res_blocks=1",",","...","attention_resolutions=","[","1","]","...",")",">",">",">","x","=","torch.randn","(","4",",","1",",","2",",","4",")",">",">",">","ts","=","torch.tensor","(","[","10",",","100",",","50",",","25","]",")",">",">",">","out","=","model","(","x",",","ts",")",">",">",">","out.shape","torch.Size","(","[","4",",","3",",","16",",","32","]",")"]
docstr: Apply the model to an input batch.
  
  Arguments
  ---------
  x:  torch.Tensor
      an [N x C x ...] Tensor of inputs.
  timesteps: torch.Tensor
      a 1-D batch of timesteps.
  Returns
  --------
  result: torch.Tensor
      an [N x K] Tensor of outputs.
docstr_tok: ["Apply","the","model","to","an","input","batch",".","Arguments","--","--","--","--","-","x",":","torch.Tensor","an","[","N","x","C","x","...","]","Tensor","of","inputs",".","timesteps",":","torch.Tensor","a","1-D","batch","of","timesteps",".","Returns","--","--","--","--","result",":","torch.Tensor","an","[","N","x","K","]","Tensor","of","outputs","."]
code: def forward(self, x, timesteps=None):
      """
      Apply the model to an input batch.
  
      Arguments
      ---------
      x:  torch.Tensor
          an [N x C x ...] Tensor of inputs.
      timesteps: torch.Tensor
          a 1-D batch of timesteps.
      Returns
      --------
      result: torch.Tensor
          an [N x K] Tensor of outputs.
      """
      emb = None
      if timesteps is not None:
          emb = self.time_embed(
              timestep_embedding(timesteps, self.model_channels)
          )
  
      h = x.type(self.dtype)
      h = self.input_block(h, emb)
      h = self.middle_block(h, emb)
      for module in self.upsample_blocks:
          h = module(h, emb)
      h = self.out(h)
      return h
  
imports: from abc import abstractmethod
  from speechbrain.utils.data_utils import pad_divisible
  from .autoencoders import NormalizingAutoencoder
  import math
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the DecoderUNetModel obj's underlying member method code with example usage: obj.forward(x, timesteps) and returns: h