acc_train = {}
acc_val = {}
# Iterate over epochs
for epoch in range(epochs):

    n_correct=0; n_samples=0; n_true_OK=0
    for idx, (images, labels) in enumerate(train_loader):
        model.train()
        # Push data to gpu if available
        images, labels = images.to(device), labels.to(device)

        # Forward pass
        outputs = model(images)
        l = loss(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        l.backward()
        optimizer.step()

        # Get prediced labels (.max returns (value,index))
        _, y_pred = torch.max(outputs.data, 1)

        # Count correct classifications
        n_correct += (y_pred == labels).sum().item()
        n_true_OK += (labels == 1).sum().item()
        n_samples += labels.size(0)

    # At end of epoch: Eval accuracy and print information
    if (epoch+1) % 2 == 0:
        model.eval()
        # Calculate accuracy
        acc_train[epoch+1] = n_correct / n_samples
        true_OK = n_true_OK / n_samples
        acc_val[epoch+1] = val_test(val_loader, model)[0]

        # Print info
        print (f"Epoch [{epoch+1}/{epochs}], Loss: {l.item():.4f}")
        print(f"      Training accuracy: {acc_train[epoch+1]*100:.2f}%")
        print(f"      True OK: {true_OK*100:.3f}%")
        print(f"      Validation accuracy: {acc_val[epoch+1]*100:.2f}%")

# Save model and state_dict
torch.save(model, "model.pth")