    def forward(self, x, mel, softmax=False):
        """Forward step

        Args:
            x (Tensor): One-hot encoded audio signal, shape (B x C x T)
            mel (Tensor): Local conditioning features,
              shape (B x cin_channels x T)
            softmax (bool): Whether applies softmax or not.

        Returns:
            Tensor: output, shape B x out_channels x T
        """
        B, _, T = x.size()

        mel = self.upsample_net(mel)
        assert mel.shape[-1] == x.shape[-1]

        x = self.first_conv(x)
        skips = 0
        for f in self.conv_layers:
            x, h = f(x, mel)
            skips += h
        skips *= math.sqrt(1.0 / len(self.conv_layers))

        x = skips
        for f in self.last_conv_layers:
            x = f(x)

        x = F.softmax(x, dim=1) if softmax else x

        return x
Forward step

Args:
x (Tensor): One-hot encoded audio signal, shape (B x C x T)
mel (Tensor): Local conditioning features,
shape (B x cin_channels x T)
softmax (bool): Whether applies softmax or not.

Returns:
Tensor: output, shape B x out_channels x T
models/vocoders/autoregressive/wavenet/wavenet.py