First of all, thank you both for all of your help, you’ve been really helpful and i understand so much more now.
I think the problem i’m facing now is about the way the script i based the model on calculates and represents the accuracy.
Here’s the code that gets the accuracy and loss:
def train_model(model, optimizer, loss_fn, train_loader, val_loader, epochs):
    train_accuracies, train_losses, val_accuracies, val_losses = [], [], [], []
    val_loss = AverageMeter()
    val_accuracy = AverageMeter()
    train_loss = AverageMeter()
    train_accuracy = AverageMeter()

    for epoch in range(epochs):
        # train
        model.train()
        train_loss.reset()
        train_accuracy.reset()
        train_loop = tqdm(train_loader, unit=" batches")  # For printing the progress bar
        for data, target in train_loop:
            train_loop.set_description('[TRAIN] Epoch {}/{}'.format(epoch + 1, epochs))
            data, target = data.float().to(device), target.long().to(device)
            target = target.unsqueeze(-1).long()
            optimizer.zero_grad()
            output = model(data)
            loss = loss_fn(output, torch.max(target, 1)[1])
            loss.backward()
            optimizer.step()
            train_loss.update(loss.item(), n=len(target))
            pred = output.round()  # get the prediction
            acc = pred.eq(target.view_as(pred)).sum().item() / len(target.long())
            train_accuracy.update(acc, n=len(target))
            train_loop.set_postfix(loss=train_loss.avg, accuracy=train_accuracy.avg)

        train_losses.append(train_loss.avg)
        train_accuracies.append(train_accuracy.avg)

        # validation
        model.eval()
        val_loss.reset()
        val_accuracy.reset()
        val_loop = tqdm(val_loader, unit=" batches")  # For printing the progress bar
        with torch.no_grad():
            for data, target in val_loop:
                val_loop.set_description('[VAL] Epoch {}/{}'.format(epoch + 1, epochs))
                data, target = data.float().to(device), target.long().to(device)
                target = target.unsqueeze(-1).long()
                output = model(data)
                loss = loss_fn(output, torch.max(target, 1)[1])
                val_loss.update(loss.item(), n=len(target))
                pred = output.round()  # get the prediction
                acc = pred.eq(target.view_as(pred)).sum().item() / len(target.long())
                val_accuracy.update(acc, n=len(target))
                val_loop.set_postfix(loss=val_loss.avg, accuracy=val_accuracy.avg)

        val_losses.append(val_loss.avg)
        val_accuracies.append(val_accuracy.avg)

    return train_accuracies, train_losses, val_accuracies, val_losses


epochs = 30
train_accuracies, train_losses, val_accuracies, val_losses = train_model(model, optimizer, loss_fn, train_loader, val_loader, epochs)

torch.save(model.state_dict(), 'Types_of_lane_small_3_class.pt')
np.savez('history_Types_of_lane_small_3_class.npz', train_accuracies=train_accuracies, train_losses=train_losses, val_accuracies=val_accuracies, val_losses=val_losses)

epochs = range(len(train_accuracies))

plt.plot(epochs, train_accuracies, 'b', label='Training acc')
plt.plot(epochs, val_accuracies, 'r', label='Validation acc')
plt.title('Training and validation accuracy')
plt.legend()

plt.figure()

plt.plot(epochs, train_losses, 'b', label='Training loss')
plt.plot(epochs, val_losses, 'r', label='Validation loss')
plt.title('Training and validation loss')
plt.legend()

plt.show()

And the error i’m getting now, with nn.LogSoftmax with dim=1 and the number of output features set to 3

Traceback (most recent call last):
  File "Model3class.py", line 265, in <module>
    train_accuracies, train_losses, val_accuracies, val_losses = train_model(model, optimizer, loss_fn, train_loader, val_loader, epochs)
  File "Model3class.py", line 233, in train_model
    acc = pred.eq(target.view_as(pred)).sum().item() / len(target.long())
RuntimeError: shape '[32, 3]' is invalid for input of size 32


I don’t really know if this way of getting the accuracy is valid for a 3 or 4 class classification model