Loss Functions and What They Expect

Pick a loss that matches your task — and read the input contract

Most "my model isn't learning" stories trace back to a mismatch between the loss function and what the model outputs. The fix is usually a one-line change once you've spotted it.

Classification

nn.CrossEntropyLoss — multi-class. Inputs: raw logits shape (N, C); targets: int64 class indices shape (N,). Combines log_softmax + NLL internally for numerical stability — do NOT softmax before it.
nn.BCEWithLogitsLoss — binary OR multi-label. Inputs: raw logits; targets: floats in {0, 1}. Combines sigmoid + BCE for stability — do NOT sigmoid before it.
nn.NLLLoss — multi-class but expects log-probabilities instead of logits. Use only if you've already applied log_softmax.

Regression

nn.MSELoss — mean squared error. The default for regression. Penalizes outliers heavily because of the square.
nn.L1Loss — mean absolute error. More robust to outliers.
nn.SmoothL1Loss / nn.HuberLoss — L2 near zero, L1 in the tails. The "best of both" for noisy regression.

Less common but useful

nn.KLDivLoss — KL divergence between two distributions. Used in knowledge distillation.
nn.CosineEmbeddingLoss — for similarity-based learning (face verification, embedding similarity).
nn.TripletMarginLoss — for metric learning where you have anchor/positive/negative triples.

Code

Multi-class classification — CrossEntropyLoss·python

import torch
import torch.nn as nn

# Model outputs raw logits, NOT softmax probabilities
model = nn.Linear(10, 5)               # 5 classes
criterion = nn.CrossEntropyLoss()

x = torch.randn(16, 10)
targets = torch.randint(0, 5, (16,))   # int64 class indices

logits = model(x)
loss = criterion(logits, targets)
print(loss.item())                      # scalar

Binary / multi-label — BCEWithLogitsLoss·python

import torch
import torch.nn as nn

# Binary classifier — single logit per sample
model_bin = nn.Linear(10, 1)
criterion = nn.BCEWithLogitsLoss()

x = torch.randn(16, 10)
y_bin = torch.randint(0, 2, (16, 1)).float()   # MUST be float, not int
loss = criterion(model_bin(x), y_bin)

# Multi-label — N binary outputs per sample
model_ml = nn.Linear(10, 5)             # 5 independent binary labels
y_ml = torch.randint(0, 2, (16, 5)).float()
loss_ml = criterion(model_ml(x), y_ml)
print(loss_ml.item())

Regression — pick by outlier behavior·python

import torch
import torch.nn as nn

x = torch.randn(16, 10)
y = torch.randn(16, 2)
model = nn.Linear(10, 2)

mse = nn.MSELoss()                      # default; squared, sensitive to outliers
l1  = nn.L1Loss()                       # robust, less sharp gradient near zero
huber = nn.SmoothL1Loss(beta=1.0)       # L2 near zero, L1 in tails

for name, fn in [('MSE', mse), ('L1', l1), ('Huber', huber)]:
    print(f"{name}: {fn(model(x), y).item():.4f}")

class_weight — handle imbalanced classes·python

import torch
import torch.nn as nn

# Suppose class 0 is 9x more common than class 1
# Heavier weight on the rare class so the loss cares more about it
weights = torch.tensor([1.0, 9.0])      # one weight per class
criterion = nn.CrossEntropyLoss(weight=weights)

logits = torch.randn(16, 2)
targets = torch.cat([torch.zeros(14, dtype=torch.long), torch.ones(2, dtype=torch.long)])
loss = criterion(logits, targets)
print(loss.item())

Exercise

Build a classifier with output shape (B, 5). Apply BOTH softmax-then-NLLLoss AND raw-logits-with-CrossEntropyLoss on the same input + target. Verify both produce the same loss value (they're equivalent), then verify that softmax-then-CrossEntropyLoss produces the wrong (smaller) value because of the double softmax.