Are SWA BN stats being updated before val_loop()?

[chanshing]

Is the StochasticWeightAveraging implementation updating the BN stats before a validation starts?

I had a look and it doesn't seem so, or I may have missed it. It seems to only update the BN stats at the very end of all training, but we also need to update them for validation, as shown in the original implementation: https://github.com/timgaripov/swa/blob/411b2fcad59bec60c6c9eb1eb19ab906540e5ea2/train.py#L143-L161

for epoch in range(start_epoch, args.epochs):
    time_ep = time.time()


    lr = schedule(epoch)
    utils.adjust_learning_rate(optimizer, lr)
    train_res = utils.train_epoch(loaders['train'], model, criterion, optimizer)
    if epoch == 0 or epoch % args.eval_freq == args.eval_freq - 1 or epoch == args.epochs - 1:
        test_res = utils.eval(loaders['test'], model, criterion)
    else:
        test_res = {'loss': None, 'accuracy': None}


    if args.swa and (epoch + 1) >= args.swa_start and (epoch + 1 - args.swa_start) % args.swa_c_epochs == 0:
        utils.moving_average(swa_model, model, 1.0 / (swa_n + 1))
        swa_n += 1
        if epoch == 0 or epoch % args.eval_freq == args.eval_freq - 1 or epoch == args.epochs - 1:
            utils.bn_update(loaders['train'], swa_model)
            swa_res = utils.eval(loaders['test'], swa_model, criterion)
        else:
            swa_res = {'loss': None, 'accuracy': None}

We can see in the last few lines that it is updating the swa_model BN stats right before validation.
I was hoping that someone more familiar with the lightning implementation can help clarify my concern.