Hi! Thanks for all that! What ended up working was nn.CrossEntropyLoss() along with out = out.view(out.size(0), -1) instead of my original out = out.view(1, -1).
Now… I have a new problem, although it’s slightly unrelated.
Basically, I suspect that my model isn’t properly saving and loading.
I run the training loop where I define the directory to save to:
# Train the model

freq_net = Frequentist_CNN(num_classes, num_channels).to(device)

# Loss and optimizer
# criterion = nn.BCEWithLogitsLoss()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(freq_net.parameters(), lr=0.01)

# Where to save the model:
freq_ckpt_dir = 'saved_model/COVID/model_custom/'
# File name
freq_ckpt_name = 'freq_COVID_CNN.pt'

path_freq_CNN = freq_ckpt_dir + freq_ckpt_name

# If it doesn't exist, make it.
if not os.path.exists(freq_ckpt_dir):
    os.makedirs(freq_ckpt_dir, exist_ok=True)     


total_step = len(train_loader)
loss_list = []
acc_list = []
# Train
print("Training.....")
for epoch in range(n_epochs):
    for i, (images, labels) in enumerate(train_loader):
        print("Batch: \t{}".format(i+1))
        # Run the forward pass
        #print("Images have shape: {}".format(images.shape))
        images, labels = images.to(device), labels.to(device)
        #labels = labels.unsqueeze(1)
        outputs = freq_net(images).to(device)
        #outputs = outputs.argmax(dim=1, keepdim=True)
        #print("Outputs have shape: {}".format(outputs.shape))
        #print("Labels have shape: {}".format(labels.shape))

        loss = criterion(outputs, labels)
        loss_list.append(loss.item())

        # Backprop and perform Adam optimisation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Track the accuracy
        total = labels.size(0)
        _, predicted = torch.max(outputs.data, 1)
        correct = (predicted == labels).sum().item()
        acc_list.append(correct / total)

        if (i + 1) % 4 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Accuracy: {:.2f}%'
                  .format(epoch + 1, n_epochs, i + 1, total_step, loss.item(),
                          (correct / total) * 100))
            

        # Save.
        torch.save(freq_net.state_dict(), path_freq_CNN)  
print("Done training.")

And it gets this performance after 10 epochs:
Epoch [10/10], Step [4/5], Loss: 0.8934, Accuracy: 78.12%

But now I go load it and test:
# Test the model

freq_net.load_state_dict(torch.load(path_freq_CNN))


freq_net.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = freq_net(images).to(device)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy of the model on the 60 test images: {} %'.format((correct / total) * 100))


Only to see:
Test Accuracy of the model on the 60 test images: 50.0 %

My test-set is balanced, so it seems to be outputting just one class every time.
Is the problem that I’m not saving/loading correctly? It seems my test_loader is defined correctly, so that can’t be it…
Thanks so much!