PyTorch LSTM categorical model - output to target mapping
I have a network which outputs a vector of length two. My targets are in the form of 1 or zeros, referring to two possible categories. What is the best way to get the loss - i.e. should I transform the targets, for example into a dimension 2 vector, or should I transform the output of the network, e.g. take the location of the max number as the output?
My network looks like:
class LSTMClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim
        self.lstm1 = nn.LSTM(input_dim, hidden_dim, layer_dim, batch_first=True)
        self.lstm2 = nn.LSTM(hidden_dim, hidden_dim, layer_dim, batch_first=True)
        self.fc1 = nn.Linear(hidden_dim, 32)
        self.fc2 = nn.Linear(32, 1)
        self.dropout = nn.Dropout(p=0.2)
        self.batch_normalisation1 = nn.BatchNorm1d(layer_dim)
        self.batch_normalisation2 = nn.BatchNorm1d(2)
        self.activation = nn.Softmax(dim=2)
    
    def forward(self, x):
        h0, c0 = self.init_hidden(x)
        out, (hn1, cn1) = self.lstm1(x, (h0, c0))
        out = self.dropout(out,)
        out = self.batch_normalisation1(out)
        
        h1, c1 = self.init_hidden(out)
        out, (hn2, cn2) = self.lstm2(out, (h1, c1))
        out = self.dropout(out)
        out = self.batch_normalisation1(out)
        
        h2, c2 = self.init_hidden(out)
        out, (hn3, cn3) = self.lstm2(out, (h2, c2))
        out = self.dropout(out)
        out = self.batch_normalisation1(out)
        
        out = self.fc1(out[:, -1, :])
        out = self.dropout(out)
        out = self.fc2(out)
        return out
    
    def init_hidden(self, x):
        h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim)
        c0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim)
        return [t for t in (h0, c0)]
    
    def pred(self, x):
        out = self(x)
        return out &gt; 0

An example of input to this network is:
tensor([[[0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
         [2.3597e-04, 1.1507e-02, 8.7719e-02, 6.1093e-02, 9.5556e-01],
         [2.1474e-03, 5.3805e-03, 9.6491e-02, 2.2508e-01, 8.2222e-01]]])

which has shape torch.Size([1, 3, 5]). The target is currently 1 or 0. However, the network outputs a vector such as:
tensor([[0.5293, 0.4707]], grad_fn=&lt;SoftmaxBackward&gt;)

What would be the best way to set up the loss between these target and the network output?
Update:
I can now train the model as suggested in the answers as:
model = LSTMClassifier(5, 128, 3, 1)
Epochs = 10
batch_size = 32

criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=1e-6)

for epoch in range(Epochs):
    if epoch == 0:
        accurate = 0
        for X_instance, y_instance in zip(val_x, val_y):
            if int(y_instance) == 1 and model.pred(X_instance.view(-1, 3, 5)).item():
                accurate += 1
        print(f&quot;Untrained accuracy test set: {accurate/len(val_x)}&quot;)
    print(f&quot;Epoch {epoch + 1}&quot;)
    
    for n, (X, y) in enumerate(train_batches):
        model.train()
        optimizer.zero_grad()
        y_pred = model(X)
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()

    model.eval()
    accurate = 0
    for X_instance, y_instance in zip(val_x, val_y):
        if int(y_instance) == 1 and model.pred(X_instance.view(-1, 3, 5)).item():
            accurate += 1
    print(f&quot;Accuracy test set: {accurate/len(val_x)}&quot;)


You shouldn't use any activation at the end of your network and output only a single neuron instead of two (trained with BCEWithLogitsLoss).
Below is your neural network code with commentary and removal of unnecessary parts:
class LSTMClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim
        self.lstm1 = nn.LSTM(input_dim, hidden_dim, layer_dim, batch_first=True)
        self.lstm2 = nn.LSTM(hidden_dim, hidden_dim, layer_dim, batch_first=True)
        self.fc1 = nn.Linear(hidden_dim, 32)
        # Output 1 neuron instead of two
        self.fc2 = nn.Linear(32, 1)
        # Model should not depend on batch size
        # self.batch_size = None
        # You are not using this variable
        # self.hidden = None
        self.dropout = nn.Dropout(p=0.2)
        self.batch_normalisation1 = nn.BatchNorm1d(layer_dim)
        self.batch_normalisation2 = nn.BatchNorm1d(2)

    def forward(self, x):
        # Hidden are initialized with 0 explicitly
        # h0, c0 = self.init_hidden(x)
        out, _ = self.lstm1(x)
        # No need for initial values
        # out, (hn1, cn1) = self.lstm1(x, (h0, c0))
        out = self.dropout(out)
        out = self.batch_normalisation1(out)

        # Same for all other cells you re-init with zeros, it's implicit
        out, _ = self.lstm2(out)
        out = self.dropout(out)
        out = self.batch_normalisation1(out)

        out, _ = self.lstm2(out)
        out = self.dropout(out)
        out = self.batch_normalisation1(out)

        out = self.fc1(out[:, -1, :])
        out = self.dropout(out)
        # No need for activation
        # out = F.softmax(self.fc2(out))
        out = self.fc2(out)
        return out

    # Return True (1) or False (0)
    def pred(self, x):
        return self(x) &gt; 0

I have also added pred method which transforms logits into targets (e.g. to use with some metrics).
Basically, if your logit is lower than 0 it is False, otherwise it is True. No need for activation in this case.
