import torch
import torch.nn as nn
import torch.nn.functional as F

# a standard training step, the pattern repeated (with minor
# variations) in every PyTorch example from Chapter 16 onward
model = nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

x = torch.randn(32, 10)
y = torch.randint(0, 2, (32,))

logits = model(x)
loss = F.cross_entropy(logits, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

# switching between training and evaluation mode (affects
# dropout and batch normalization layers, Chapter 18)
model.eval()
with torch.no_grad():
    predictions = model(x).argmax(dim=1)
model.train()
