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)
```
Next, we explain each component of torch.optim.swa_utils in detail.
AveragedModel class serves to compute the weights of the SWA model. You can create an averaged model by running  swa_model = AveragedModel(model). You can then update the parameters of the averaged model by swa_model.update_parameters(model). By default, AveragedModel computes a running equal average of the parameters that you provide, but you can also use custom averaging functions with the avg_fn parameter. In the following example, ema_model computes an exponential moving average.
```python