pytorch_training.py

for epoch in range(n_epochs):
  for data, targets in trainloader:
    # Generate predictions
    out = model(data)
    # Calculate loss
    loss = criterion(out, targets)
    # Backpropagation
    loss.backward()
    # Update model parameters
    optimizer.step()