docs: add docs for `wrap_model` in TFKerasTrialContext
@@ -154,7 +154,7 @@ class TFKerasContext:
This should be used to wrap ``tf.data.Dataset`` objects immediately after
they have been created. Users should use the output of this wrapper as the
new instance of their dataset. If users create multiple datasets (e.g.,
- one for training and one for testing) users should wrap each dataset
+ one for training and one for testing), users should wrap each dataset
independently.
"""
hvd.require_horovod_type("tensorflow.keras", "TFKerasContext.wrap_dataset was called.")
@@ -173,6 +173,12 @@ class TFKerasTrialContext(det.TrialContext, TFKerasContext):
TFKerasContext.__init__(self, hvd_config)
def wrap_model(self, model: Any) -> Any:
+ """
+ This should be used to wrap ``tf.keras.Model`` objects immediately after
+ they have been created but before they have been compiled. This function
+ takes a ``tf.keras.Model`` and returns a wrapped version of the model;
+ the return value should be used in place of the original model.
+ """
return self._wrap_model_with_train_fn(model, None)
