EMA

class torch_uncertainty.models.EMA(model, momentum)

Exponential Moving Average.

Parameters:

• model (nn.Module) – The model to train and ensemble.

• momentum (float) – The momentum of the moving average.

update_wrapper(epoch=None)

Update the EMA model.

Parameters:

epoch (int) – The current epoch. For API consistency.