def __init__(self, model_fn, name):
    """Inits the module.

    Args:
      model_fn: callable that, when called with no arguments, returns a callable
          similar to a Sonnet module instance.
      name: Name for the wrapper module.

    """
    super(WrappedModelFnModule, self).__init__(name=name)
    with self._enter_variable_scope():
      self._model = model_fn()
Inits the module.

    Args:
      model_fn: callable that, when called with no arguments, returns a callable
          similar to a Sonnet module instance.
      name: Name for the wrapper module.

    
graph_nets/_base.py