scrapegoat/Neural-Audio-Codec

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7"""LSTM layers module."""8 9from torch import nn10 11 12class SLSTM(nn.Module):13    """14    LSTM without worrying about the hidden state, nor the layout of the data.15    Expects input as convolutional layout.16    """17    def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):18        super().__init__()19        self.skip = skip20        self.lstm = nn.LSTM(dimension, dimension, num_layers)21 22    def forward(self, x):23        x = x.permute(2, 0, 1)24        y, _ = self.lstm(x)25        if self.skip:26            y = y + x27        y = y.permute(1, 2, 0)28        return y29