code: def load_model(
      bento_model: str | Tag | bentoml.Model,
      device_name: str = "/device:CPU:0",
  ) -> "tf_ext.KerasModel":
      """
      Load a model from BentoML local modelstore with given name.
  
      Args:
          bento_model (``str`` ``|`` :obj:`~bentoml.Tag` ``|`` :obj:`~bentoml.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 (``str`` | ``None``):
              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:`keras.Model`: an instance of users :obj:`keras.Model` from BentoML modelstore.
  
      Examples:
  
      .. code-block:: python
  
          import bentoml
  
          # load a model back into memory:
          loaded = bentoml.keras.load_model("keras_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):
          return keras.models.load_model(
              bento_model.path,
              custom_objects=bento_model.custom_objects,
          )
  
imports: from __future__ import annotations
  import functools
  import logging
  import typing
  from types import ModuleType
  from typing import TYPE_CHECKING
  import attr
  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.utils import Params
  from ..types import LazyType
  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: keras