code: def load_model(
      bento_model: str | Tag | bentoml.Model,
      device_name: str = "/device:CPU:0",
  ) -> tf_ext.AutoTrackable | tf_ext.Module:
      """
      Load a tensorflow model from BentoML local modelstore with given name.
  
      Args:
          bento_model: Either the tag of the model to get from the store, or a
                       BentoML `~bentoml.Model` instance to load the model from.
          device_name: The device id to load the model on. The device id format
                       should be compatible with `tf.device <https://www.tensorflow.org/api_docs/python/tf/device>`_
  
      Returns:
          :obj:`SavedModel`: an instance of :obj:`SavedModel` format from BentoML modelstore.
  
      Examples:
  
      .. code-block:: python
  
          import bentoml
  
          # load a model back into memory
          model = bentoml.tensorflow.load_model("my_tensorflow_model")
  
      """  # noqa
      if not isinstance(bento_model, bentoml.Model):
          bento_model = get(bento_model)
  
      if bento_model.info.module not in (MODULE_NAME, __name__):
          raise NotFound(
              f"Model {bento_model.tag} was saved with module {bento_model.info.module}, not loading with {MODULE_NAME}."
          )
  
      if "GPU" in device_name:
          physical_devices = tf.config.list_physical_devices("GPU")
          try:
              # an optimization for GPU memory growth. But it will raise an error if any
              # tensorflow session is already created. That happens when users test runners
              # in a notebook or Python interactive shell. Thus we just ignore the error.
              tf.config.experimental.set_memory_growth(
                  physical_devices[0], True
              )
          except RuntimeError:
              pass
  
      with tf.device(device_name):  # type: ignore (tf.device is a context manager)
          tf_model: tf_ext.AutoTrackable = tf.saved_model.load(
              bento_model.path
          )
          return tf_model
  
imports: from __future__ import annotations
  import contextlib
  import itertools
  import logging
  import pickle
  import typing
  from types import ModuleType
  from typing import TYPE_CHECKING
  import bentoml
  from bentoml import Runnable
  from bentoml import Tag
  from bentoml.exceptions import MissingDependencyException
  from bentoml.exceptions import NotFound
  from bentoml.models import ModelContext
  from ..models.model import ModelSignature
  from ..models.model import PartialKwargsModelOptions
  from ..runner.container import DataContainer
  from ..runner.container import DataContainerRegistry
  from ..runner.container import Payload
  from ..types import LazyType
  from .utils.tensorflow import cast_py_args_to_tf_function_args
  from .utils.tensorflow import get_input_signatures_v2
  from .utils.tensorflow import get_output_signatures_v2
  from .utils.tensorflow import get_restorable_functions
  from .utils.tensorflow import get_tf_version
label_desc: Use image node_id 2 for calling a global function with example usage: load_model(bento_model, device_name) and returns: tf_model