Sure! So in SWA two models are maintained: the model and the swa_model. The latter is the averaged model. What we “train” is model (we backprop this, update its weights, etc.), and only every now and then we update swa_model with model by averaging. That’s why the batchnorm stats in swa_model needs separate updating. From the Pytorch website:

One important detail is the batch normalization. Batch normalization layers compute running statistics of activations during training. Note that the SWA averages of the weights are never used to make predictions during training. So the batch normalization layers do not have the activation statistics computed at the end of training. We can compute these statistics by doing a single forward pass on the train data with the SWA model.

And here’s the code snippet provided by Pytorch, where we see BN stats of swa_model being updated at the end right before testing:
for epoch in range(100):
      for input, target in loader:
          optimizer.zero_grad()
          loss_fn(model(input), target).backward()
          optimizer.step()
      if epoch > swa_start:
          swa_model.update_parameters(model)
          swa_scheduler.step()
      else:
          scheduler.step()

# Update bn statistics for the swa_model at the end
torch.optim.swa_utils.update_bn(loader, swa_model)
# Use swa_model to make predictions on test data 
preds = swa_model(test_input)

The code snippet doesn’t show the validation step. So my question was whether we should also update_bn() the swa_model for validating it.