import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


features = torch.tensor(
    [
        [0.0, 0.1],
        [0.2, 0.0],
        [1.0, 1.1],
        [1.2, 0.9],
    ],
    dtype=torch.float32,
)
labels = torch.tensor([0, 0, 1, 1])
validation_data = TensorDataset(features, labels)
validation_loader = DataLoader(validation_data, batch_size=2)

model = nn.Sequential(
    nn.Linear(2, 2),
    nn.Dropout(p=0.75),
)

with torch.no_grad():
    model[0].weight.copy_(torch.tensor([[1.0, -1.0], [1.0, 1.0]]))
    model[0].bias.copy_(torch.tensor([0.25, -1.0]))

loss_fn = nn.CrossEntropyLoss(reduction="sum")

model.eval()
loss_total = 0.0
correct = 0
sample_count = 0
gradient_flags = []

with torch.inference_mode():
    for batch_features, batch_labels in validation_loader:
        logits = model(batch_features)
        loss = loss_fn(logits, batch_labels)

        loss_total += loss.item()
        predictions = logits.argmax(dim=1)
        correct += (predictions == batch_labels).sum().item()
        sample_count += batch_labels.size(0)
        gradient_flags.append(logits.requires_grad)

average_loss = loss_total / sample_count
accuracy = correct / sample_count

print(f"model training mode: {model.training}")
print(f"dropout training mode: {model[1].training}")
print(f"logits require grad: {any(gradient_flags)}")
print(f"validation loss: {average_loss:.4f}")
print(f"validation accuracy: {correct}/{sample_count} ({accuracy:.0%})")