Tensors, Autograd & Training Loop

Tensors & Devices

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
y = torch.randn(2, 2)
z = x + y  # broadcasting applies when shapes are compatible but not identical

# Move to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = x.to(device)

# Reshape without copying data (when contiguous)
flat = x.view(-1)      # or x.reshape(-1) — handles non-contiguous cases too

# NumPy interop — often shares underlying memory on CPU
import numpy as np
arr = x.cpu().numpy()
back = torch.from_numpy(arr)

Autograd

w = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

x = torch.tensor(3.0)
y = w * x + b
y.backward()  # computes dy/dw and dy/db

print(w.grad)  # tensor(3.) — dy/dw = x
print(b.grad)  # tensor(1.) — dy/db = 1

# Gradients ACCUMULATE across .backward() calls — clear them each step
w.grad.zero_()

# Disable tracking for inference — saves memory, no computation graph built
with torch.no_grad():
    prediction = w * x + b

# In-place ops on tensors needed for the backward pass can break autograd:
# x += 1        # risky if x's original value is needed later
# x = x + 1     # safe — creates a new tensor instead

Model & Training Loop

import torch.nn as nn
import torch.optim as optim

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = MLP().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()  # expects raw logits, not softmaxed probs

for epoch in range(num_epochs):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()          # clear accumulated gradients
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()                # compute gradients
        optimizer.step()               # apply the update

    # Validation — no_grad + eval mode
    model.eval()
    with torch.no_grad():
        for images, labels in val_loader:
            preds = model(images.to(device))
            # accumulate metrics with .item() — NOT the raw tensor,
            # which would keep the whole computation graph alive