code: def set_training(model, mode):
      
      if mode is None:
          yield
          return
      old_mode = model.training
      if old_mode != mode:
          model.train(mode)
      try:
          yield
      finally:
          if old_mode != mode:
              model.train(old_mode)
docs: 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.