import torch
import torch.nn as nn
import torch.nn.functional as F

class AudioExpert(nn.Module):
    def __init__(self):
        super(AudioExpert, self).__init__()
        self.rnn = nn.LSTM(input_size=40, hidden_size=128, num_layers=2, batch_first=True)
        self.fc1 = nn.Linear(128, 128)

    def forward(self, x):
        h0 = torch.zeros(2, x.size(0), 128).to(x.device)
        c0 = torch.zeros(2, x.size(0), 128).to(x.device)
        x, _ = self.rnn(x, (h0, c0))
        x = F.relu(self.fc1(x[:, -1, :]))
        return x