Eric

import torch, torch.nn as nn, torch.optim as optim
model = nn.Sequential(
    nn.Linear(784, 256), nn.BatchNorm1d(256), nn.ReLU(),
    nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(),
    nn.Linear(128, 10)
)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
# Initialize weights — Born with certain talents (random init)
for epoch in range(100):
    # Forward pass — Understanding the world with existing knowledge
    predictions = model(train_images)
    # Compute loss — Recognizing our shortcomings
    loss = criterion(predictions, train_labels)
    # Backpropagation — Learning from mistakes, computing gradients
    loss.backward()
    # Update weights — Growth and change, fine-tuning along gradients
    optimizer.step()
    optimizer.zero_grad()
# Ultimately, we learn not to replicate the past, but to generalize to the future