I’m building a multiclass classification model using a GRU. I’m struggling to get my head around how to shape the output such that the loss can be calculated.
Here is my code:
class Network(nn.Module):
def __init__(self, input_dim, embedding_dim, hidden_dim, output_dim, weights=None):
    super(Net, self).__init__()
    self.embedding = nn.Embedding(input_dim, embedding_dim)
    if weights is not None:
        self.embedding.weight.data.copy_(weights)

    self.embedding.weight.requires_grad = False
    self.lstm = nn.LSTM(embedding_dim, hidden_dim)
    self.linear = nn.Linear(hidden_dim, output_dim)
    self.sigmoid = nn.Sigmoid()

def forward(self, inp):
    embed_out = self.embedding(inp)
    out, (hidden1, hidden2) = self.lstm(embed_out)
    out = self.sigmoid(self.linear(hidden1.squeeze(0)))
    return out.squeeze(-1)

The output shape in my batch [342, 51] whereas the shape of the label is [32, 51] which means the loss can’t be calculated. Any pointers would be very helpful - thank you.