PyTorch MSELoss: reduction, F.mse_loss and Shapes

nn.MSELoss is the mean squared error loss in PyTorch. It squares the gap between each prediction and its target, then averages.

criterion = nn.MSELoss()
loss = criterion(predictions, targets)

There’s a functional twin, F.mse_loss, that does exactly the same arithmetic without the object. Which you use is a style choice, and I’ll show why below.

Runs and screenshots come from PyTorch 2.14.0, Python 3.12.5.

What nn.MSELoss actually computes

Nothing is hidden here, so it’s worth seeing the arithmetic next to the call:

import torch
import torch.nn as nn

criterion = nn.MSELoss()

predictions = torch.tensor([2.5, 0.0, 2.1, 7.8])
targets     = torch.tensor([3.0, -0.5, 2.0, 7.0])

loss = criterion(predictions, targets)
print("nn.MSELoss :", loss.item())

# the same calculation written out
squared = (predictions - targets) ** 2
print("differences:", (predictions - targets).tolist())
print("squared    :", squared.tolist())
print("their mean :", squared.mean().item())

Output:

nn.MSELoss : 0.2875000834465027
differences: [-0.5, 0.5, 0.09999990463256836, 0.8000001907348633]
squared    : [0.25, 0.25, 0.00999998115003109, 0.6400002837181091]
their mean : 0.2875000834465027

Square the differences, take the mean. Squaring is what makes MSE punish a big miss far more than several small ones, which is the whole reason you’d pick it.

MSELoss reduction: mean, sum and none

The reduction argument decides what comes back. It’s the argument people get wrong most often, because two of the three don’t return a single number you can call .backward() on without more work:

import torch
import torch.nn as nn

predictions = torch.tensor([2.5, 0.0, 2.1, 7.8])
targets     = torch.tensor([3.0, -0.5, 2.0, 7.0])

for how in ("mean", "sum", "none"):
    loss = nn.MSELoss(reduction=how)(predictions, targets)
    print(f"reduction={how:<6} -> {loss}")

Output:

reduction=mean   -> 0.2875000834465027
reduction=sum    -> 1.1500003337860107
reduction=none   -> tensor([0.2500, 0.2500, 0.0100, 0.6400])

reduction	Returns	Use it when
'mean'	One number, the default	Almost always
'sum'	One number, not divided	You’re averaging over batches yourself
'none'	A tensor the shape of the input	You want to weight or mask per sample

'none' is the one worth remembering. It’s how you build a weighted loss, and there’s an example further down.

nn.MSELoss or F.mse_loss?

Both run the same kernel and return the same number:

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

predictions = torch.tensor([2.5, 0.0, 2.1, 7.8])
targets     = torch.tensor([3.0, -0.5, 2.0, 7.0])

from_module = nn.MSELoss()(predictions, targets)
from_function = F.mse_loss(predictions, targets)

print("nn.MSELoss :", from_module.item())
print("F.mse_loss :", from_function.item())
print("identical  :", torch.equal(from_module, from_function))

# the module holds no weights at all, so neither choice affects training
print("MSELoss parameters:", list(nn.MSELoss().parameters()))
print("F.mse_loss with reduction='sum':", F.mse_loss(predictions, targets, reduction="sum").item())

Output:

nn.MSELoss : 0.2875000834465027
F.mse_loss : 0.2875000834465027
identical  : True
MSELoss parameters: []
F.mse_loss with reduction='sum': 1.1500003337860107

The revealing line is the empty parameter list. nn.MSELoss holds no learnable weights, so unlike a layer, the module form buys you nothing at training time.

Use the module when you want to configure the loss once in __init__ and forget it, which reads well beside your nn.Linear layers. Use the function for a one-off call or inside a custom loss.

The MSELoss target size warning

This is the MSELoss bug I see most, and it doesn’t raise an error. A model outputs (N, 1), the dataset gives targets of shape (N,), and PyTorch broadcasts them into an (N, N) grid.

You get a warning and a number that looks plausible:

import warnings
import torch
import torch.nn as nn

predictions = torch.tensor([[1.0], [2.0], [3.0]])    # shape (3, 1), straight from a Linear layer
targets = torch.tensor([1.5, 2.5, 3.5])              # shape (3,), straight from the dataset

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    wrong = nn.MSELoss()(predictions, targets)
    for warning in caught:
        print("WARNING:", warning.message)

print("shapes      :", tuple(predictions.shape), "vs", tuple(targets.shape))
print("broadcast to:", tuple((predictions - targets).shape), "-> 9 comparisons, not 3")
print("wrong loss  :", round(wrong.item(), 6))

fixed = nn.MSELoss()(predictions.squeeze(1), targets)
print("right loss  :", round(fixed.item(), 6))

Output:

WARNING: Using a target size (torch.Size([3])) that is different to the input size (torch.Size([3, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
shapes      : (3, 1) vs (3,)
broadcast to: (3, 3) -> 9 comparisons, not 3
wrong loss  : 1.583333
right loss  : 0.25

The two losses aren’t close, and nothing crashed. If your loss plateaus at a strange value, check the shapes before you touch the learning rate.

Fix it with .squeeze(1) on the output or .unsqueeze(1) on the target. Whichever you pick, make both sides the same shape before the loss sees them.

Weighted MSELoss with reduction=’none’

Set reduction='none', multiply by your weights, then reduce by hand:

import torch
import torch.nn as nn

predictions = torch.tensor([2.0, 4.0, 6.0])
targets     = torch.tensor([1.0, 4.0, 10.0])
weights     = torch.tensor([1.0, 1.0, 5.0])     # the third sample matters five times as much

per_element = nn.MSELoss(reduction="none")(predictions, targets)
print("per element:", per_element.tolist())
print("plain mean :", per_element.mean().item())
print("weighted   :", ((per_element * weights).sum() / weights.sum()).item())

Output:

per element: [1.0, 0.0, 16.0]
plain mean : 5.666666507720947
weighted   : 11.571428298950195

Dividing by weights.sum() rather than the count keeps it a weighted mean. Skip that and the loss scale moves whenever the weights do, which makes learning rates hard to reason about.

MSELoss, L1Loss and the Huber family

One outlier is enough to show why these exist:

import torch
import torch.nn as nn

predictions = torch.tensor([1.0, 2.0, 3.0, 4.0])
targets     = torch.tensor([1.1, 2.1, 2.9, 12.0])    # the last one is an outlier

losses = (("MSELoss", nn.MSELoss()), ("L1Loss", nn.L1Loss()),
          ("SmoothL1Loss", nn.SmoothL1Loss()), ("HuberLoss", nn.HuberLoss()))

for name, criterion in losses:
    print(f"{name:<13} {criterion(predictions, targets).item():8.4f}")

Output:

MSELoss        16.0075
L1Loss          2.0750
SmoothL1Loss    1.8787
HuberLoss       1.8787

MSE squares that 8.0 error into 64, which drags the mean to 16.01. L1 leaves it at 8 and lands on 2.08.

Smooth L1 and Huber sit between the two: quadratic for small errors, linear for large ones, which is why detection models reach for them.

Those last two print the same number here because their defaults line up — beta=1.0 and delta=1.0 describe the same curve. Change either and they part company.

So: MSE when large errors genuinely matter more, L1 or Huber when your data has outliers you don’t want dominating the gradient.

MSELoss in a real training loop

Here it is doing its job, recovering a line from noisy data:

import torch
import torch.nn as nn

torch.manual_seed(0)
x = torch.randn(200, 1)
y = 3 * x + 2 + 0.1 * torch.randn(200, 1)      # the line we hope to recover

model = nn.Linear(1, 1)
criterion = nn.MSELoss()
optimiser = torch.optim.SGD(model.parameters(), lr=0.1)

for epoch in range(1, 101):
    optimiser.zero_grad()
    loss = criterion(model(x), y)              # both sides are (200, 1): no broadcasting
    loss.backward()
    optimiser.step()
    if epoch == 1 or epoch % 25 == 0:
        print(f"epoch {epoch:>3}   loss {loss.item():.6f}")

print(f"learned  y = {model.weight.item():.3f}x + {model.bias.item():.3f}")
print("true     y = 3.000x + 2.000")

Output:

epoch   1   loss 7.829763
epoch  25   loss 0.011619
epoch  50   loss 0.011244
epoch  75   loss 0.011244
epoch 100   loss 0.011244
learned  y = 2.992x + 1.996
true     y = 3.000x + 2.000

Both sides of the loss are (200, 1) here, which is exactly the shape discipline the warning above is asking for. The same care applies when you join batches with torch.cat.

More PyTorch guides:

• PyTorch nn.Linear
• PyTorch nn.Conv2d
• torch.cat for joining tensors
• torch.stack in PyTorch
• PyTorch nn.Conv1d

Frequently asked questions