Inputs on wildly different scales make training slow or unstable, and inside a network the outputs of hidden layers drift in scale as the weights change. A batch normalisation layer standardises each feature over the batch, to mean 0 and standard deviation 1, then applies a learnable scale and shift.

import torch
from torch import nn

bn = nn.BatchNorm1d(2)
x = torch.tensor([[1.0, 10.0],
                  [3.0, 30.0]])
print(bn(x))  # each column: -1 and 1
print(bn.running_mean)  # [0.2, 2.0]
bn.eval()
print(bn(x[:1]))  # one row is fine now

In training mode the layer uses the statistics of the current batch, and moves running averages of them 10% of the way toward each batch. In eval mode it uses the running averages instead, so a prediction does not depend on what else is in the batch. A single row has nothing to be standardised against, so in training mode it raises an error.

Write make_model(features) returning an nn.Sequential of nn.BatchNorm1d(features) and nn.Linear(features, 1), and train(model, x, y, steps=200) that trains it in training mode with torch.optim.SGD at lr=0.1 on the mean squared error and returns it. Write predict_one(model, row) returning the prediction for one 1-D row as a float.

The explanation opens here once every check passes.