mighty.utils.prepare.ModelMode

class mighty.utils.prepare.ModelMode(mode, requires_grad)

Stores the model state with its parameters to be restored later on.

Parameters:

modebool

Original model mode extracted as model.training.

requires_graddict

A dict with keys that match model.named_parameters() dict which store a boolean state of each model parameter.

Methods

__init__(mode, requires_grad)
restore(model)	Restore the original state of the model and its parameters.

restore(model)

Restore the original state of the model and its parameters.

Parameters:

modelnn.Module

A model that was used as the input to prepare_eval() function.