class_docstr: Base context class that contains runtime information for any Determined
  workflow that uses the ``tf.keras`` API.
class_docstr_tok: ["Base","context","class","that","contains","runtime","information","for","any","Determined","workflow","that","uses","the","``","tf.keras","``","API","."]
docstr: 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.
  
  Args:
      model: tf.keras.Model
docstr_tok: ["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",".","Args",":","model",":","tf.keras.Model"]
code: 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.
  
      Args:
          model: tf.keras.Model
      """
  
      if not self.env.managed_training:
          return model
      return self._wrap_model_with_train_fn(model, None)
  
imports: import inspect
  import logging
  from typing import Any, Callable, Dict, List, NamedTuple, Optional, Union, cast
  import tensorflow
  import determined
  from determined import errors, keras, util
  from determined.common import check
  from determined.horovod import hvd
label_desc: Use image node_id 11 for calling the TFKerasTrialContext obj's underlying member method code with example usage: obj.wrap_model(model) and returns: self, model