magnet

magnet.eval(*modules)

A Context Manger that makes it easy to run computations in eval mode.

It sets modules in their eval mode and ensures that gradients are not computed.

This is a more wholesome option than torch.no_grad() since many Modules (BatchNorm, Dropout etc.) behave differently while training and testing.

Examples:

>>> import magnet as mag

>>> import magnet.nodes as mn

>>> import torch

>>> model = mn.Linear(10)

>>> x = torch.randn(4, 3)

>>> # Using eval() as context manager
>>> with mag.eval(model):
>>>     model(x)

>>> # Use as decorator
>>> @mag.eval(model)
>>> def foo():
>>>     return model(x)
>>> foo()

>>> # The modules can also be given at runtime by specifying no arguments
>>> @mag.eval
>>> def foo(model):
>>>     return model(x)
>>> foo()
>>> # The method then takes modules from the arguments
>>> # to the decorated function.