torch_uncertainty.models.mc_dropout

torch_uncertainty.models.mc_dropout(model, num_estimators, last_layer=False, on_batch=True)

MC Dropout wrapper for a model.

Parameters:

• model (nn.Module) – model to wrap

• num_estimators (int) – number of estimators to use

• last_layer (bool, optional) – whether to apply dropout to the last layer only. Defaults to False.

• on_batch (bool) – Increase the batch_size to perform MC-Dropout. Otherwise in a for loop to reduce memory footprint. Defaults to true.