@contextlib.contextmanager<EOL>def set_training(model, mode):
if mode is None:<EOL><INDENT>yield<EOL>return<EOL><DEDENT>old_mode = model.training<EOL>if old_mode != mode:<EOL><INDENT>model.train(mode)<EOL><DEDENT>try:<EOL><INDENT>yield<EOL><DEDENT>finally:<EOL><INDENT>if old_mode != mode:<EOL><INDENT>model.train(old_mode)<EOL><DEDENT><DEDENT>
A context manager to temporarily set the training mode of 'model'
to 'mode', resetting it when we exit the with-block.  A no-op if
mode is None.