zero.module.evaluation

zero.module.evaluation(*modules)

Context-manager for models evaluation.

Warning

The function must be used only as a context manager as shown below in the examples. The behaviour for call without the with keyword is unspecified.

This code…:

model.eval()
with torch.no_grad():
    ...

…is equivalent to

with evaluation(model):
    ...

Parameters

modules (torch.nn.modules.module.Module) –

See also

• ecall

Examples

a = torch.nn.Linear(1, 1)
b = torch.nn.Linear(2, 2)
with evaluation(a):
    ...
with evaluation(a, b):
    ...

model = torch.nn.Linear(1, 1)
for grad in False, True:
    for train in False, True:
        torch.set_grad_enabled(grad)
        model.train(train)
        with evaluation(model):
            assert not model.training
            assert not torch.is_grad_enabled()
            ...
        assert torch.is_grad_enabled() == grad_before_context
        # model.training is unspecified here