Spaces:
Sleeping
Sleeping
import torch.nn as nn | |
class TextClassifier(nn.Module): | |
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim): | |
super(TextClassifier, self).__init__() | |
self.embedding = nn.Embedding(vocab_size, embedding_dim) | |
self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True) | |
self.fc = nn.Linear(hidden_dim, output_dim) | |
def forward(self, x): | |
embedded = self.embedding(x) | |
lstm_out, (hidden, cell) = self.lstm(embedded) | |
last_hidden = hidden.squeeze(0) | |
logits = self.fc(last_hidden) | |
return logits |