class_docstr: WaveRNN model from *Efficient Neural Audio Synthesis* :cite:`wavernn`
  based on the implementation from `fatchord/WaveRNN <https://github.com/fatchord/WaveRNN>`_.
  
  The original implementation was introduced in *Efficient Neural Audio Synthesis*
  :cite:`kalchbrenner2018efficient`. The input channels of waveform and spectrogram have to be 1.
  The product of `upsample_scales` must equal `hop_length`.
  
  See Also:
      * `Training example <https://github.com/pytorch/audio/tree/release/0.12/examples/pipeline_wavernn>`__
      * :class:`torchaudio.pipelines.Tacotron2TTSBundle`: TTS pipeline with pretrained model.
  
  Args:
      upsample_scales: the list of upsample scales.
      n_classes: the number of output classes.
      hop_length: the number of samples between the starts of consecutive frames.
      n_res_block: the number of ResBlock in stack. (Default: ``10``)
      n_rnn: the dimension of RNN layer. (Default: ``512``)
      n_fc: the dimension of fully connected layer. (Default: ``512``)
      kernel_size: the number of kernel size in the first Conv1d layer. (Default: ``5``)
      n_freq: the number of bins in a spectrogram. (Default: ``128``)
      n_hidden: the number of hidden dimensions of resblock. (Default: ``128``)
      n_output: the number of output dimensions of melresnet. (Default: ``128``)
  
  Example
      >>> wavernn = WaveRNN(upsample_scales=[5,5,8], n_classes=512, hop_length=200)
      >>> waveform, sample_rate = torchaudio.load(file)
      >>> # waveform shape: (n_batch, n_channel, (n_time - kernel_size + 1) * hop_length)
      >>> specgram = MelSpectrogram(sample_rate)(waveform)  # shape: (n_batch, n_channel, n_freq, n_time)
      >>> output = wavernn(waveform, specgram)
      >>> # output shape: (n_batch, n_channel, (n_time - kernel_size + 1) * hop_length, n_classes)
class_docstr_tok: ["WaveRNN","model","from","*","Efficient","Neural","Audio","Synthesis","*",":","cite",":","`","wavernn","`","based","on","the","implementation","from","`","fatchord\/WaveRNN","<","https",":","\/\/github.com\/fatchord\/WaveRNN",">","`","_",".","The","original","implementation","was","introduced","in","*","Efficient","Neural","Audio","Synthesis","*",":","cite",":","`","kalchbrenner2018efficient","`",".","The","input","channels","of","waveform","and","spectrogram","have","to","be","1",".","The","product","of","`","upsample_scales","`","must","equal","`","hop_length","`",".","See","Also",":","*","`","Training","example","<","https",":","\/\/github.com\/pytorch\/audio\/tree\/release\/0.12\/examples\/pipeline_wavernn",">","`","__","*",":","class",":","`","torchaudio.pipelines.Tacotron2TTSBundle","`",":","TTS","pipeline","with","pretrained","model",".","Args",":","upsample_scales",":","the","list","of","upsample","scales",".","n_classes",":","the","number","of","output","classes",".","hop_length",":","the","number","of","samples","between","the","starts","of","consecutive","frames",".","n_res_block",":","the","number","of","ResBlock","in","stack",".","(","Default",":","``","10","``",")","n_rnn",":","the","dimension","of","RNN","layer",".","(","Default",":","``","512","``",")","n_fc",":","the","dimension","of","fully","connected","layer",".","(","Default",":","``","512","``",")","kernel_size",":","the","number","of","kernel","size","in","the","first","Conv1d","layer",".","(","Default",":","``","5","``",")","n_freq",":","the","number","of","bins","in","a","spectrogram",".","(","Default",":","``","128","``",")","n_hidden",":","the","number","of","hidden","dimensions","of","resblock",".","(","Default",":","``","128","``",")","n_output",":","the","number","of","output","dimensions","of","melresnet",".","(","Default",":","``","128","``",")","Example",">",">",">","wavernn","=","WaveRNN","(","upsample_scales=","[","5,5,8","]",",","n_classes=512",",","hop_length=200",")",">",">",">","waveform",",","sample_rate","=","torchaudio.load","(","file",")",">",">",">","#","waveform","shape",":","(","n_batch",",","n_channel",",","(","n_time","-","kernel_size","+","1",")","*","hop_length",")",">",">",">","specgram","=","MelSpectrogram","(","sample_rate",")","(","waveform",")","#","shape",":","(","n_batch",",","n_channel",",","n_freq",",","n_time",")",">",">",">","output","=","wavernn","(","waveform",",","specgram",")",">",">",">","#","output","shape",":","(","n_batch",",","n_channel",",","(","n_time","-","kernel_size","+","1",")","*","hop_length",",","n_classes",")"]
docstr: Pass the input through the WaveRNN model.
  
  Args:
      waveform: the input waveform to the WaveRNN layer (n_batch, 1, (n_time - kernel_size + 1) * hop_length)
      specgram: the input spectrogram to the WaveRNN layer (n_batch, 1, n_freq, n_time)
  
  Return:
      Tensor: shape (n_batch, 1, (n_time - kernel_size + 1) * hop_length, n_classes)
docstr_tok: ["Pass","the","input","through","the","WaveRNN","model",".","Args",":","waveform",":","the","input","waveform","to","the","WaveRNN","layer","(","n_batch",",","1",",","(","n_time","-","kernel_size","+","1",")","*","hop_length",")","specgram",":","the","input","spectrogram","to","the","WaveRNN","layer","(","n_batch",",","1",",","n_freq",",","n_time",")","Return",":","Tensor",":","shape","(","n_batch",",","1",",","(","n_time","-","kernel_size","+","1",")","*","hop_length",",","n_classes",")"]
code: def forward(self, waveform: Tensor, specgram: Tensor) -> Tensor:
      r"""Pass the input through the WaveRNN model.
  
      Args:
          waveform: the input waveform to the WaveRNN layer (n_batch, 1, (n_time - kernel_size + 1) * hop_length)
          specgram: the input spectrogram to the WaveRNN layer (n_batch, 1, n_freq, n_time)
  
      Return:
          Tensor: shape (n_batch, 1, (n_time - kernel_size + 1) * hop_length, n_classes)
      """
  
      if waveform.size(1) != 1:
          raise ValueError("Require the input channel of waveform is 1")
      if specgram.size(1) != 1:
          raise ValueError("Require the input channel of specgram is 1")
      # remove channel dimension until the end
      waveform, specgram = waveform.squeeze(1), specgram.squeeze(1)
  
      batch_size = waveform.size(0)
      h1 = torch.zeros(
          1,
          batch_size,
          self.n_rnn,
          dtype=waveform.dtype,
          device=waveform.device,
      )
      h2 = torch.zeros(
          1,
          batch_size,
          self.n_rnn,
          dtype=waveform.dtype,
          device=waveform.device,
      )
      # output of upsample:
      # specgram: (n_batch, n_freq, (n_time - kernel_size + 1) * total_scale)
      # aux: (n_batch, n_output, (n_time - kernel_size + 1) * total_scale)
      specgram, aux = self.upsample(specgram)
      specgram = specgram.transpose(1, 2)
      aux = aux.transpose(1, 2)
  
      aux_idx = [self.n_aux * i for i in range(5)]
      a1 = aux[:, :, aux_idx[0] : aux_idx[1]]
      a2 = aux[:, :, aux_idx[1] : aux_idx[2]]
      a3 = aux[:, :, aux_idx[2] : aux_idx[3]]
      a4 = aux[:, :, aux_idx[3] : aux_idx[4]]
  
      x = torch.cat([waveform.unsqueeze(-1), specgram, a1], dim=-1)
      x = self.fc(x)
      res = x
      x, _ = self.rnn1(x, h1)
  
      x = x + res
      res = x
      x = torch.cat([x, a2], dim=-1)
      x, _ = self.rnn2(x, h2)
  
      x = x + res
      x = torch.cat([x, a3], dim=-1)
      x = self.fc1(x)
      x = self.relu1(x)
  
      x = torch.cat([x, a4], dim=-1)
      x = self.fc2(x)
      x = self.relu2(x)
      x = self.fc3(x)
  
      # bring back channel dimension
      return x.unsqueeze(1)
  
imports: import math
  from typing import List, Optional, Tuple
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
label_desc: Use image node_id 2 for calling the WaveRNN obj's underlying member method code with example usage: obj.forward(waveform, specgram) and returns: x