code: def load_model(model_uri, dst_path=None, **kwargs):
      """
      Load a PyTorch model from a local file or a run.
  
      :param model_uri: The location, in URI format, of the MLflow model, for example:
  
          - ``/Users/me/path/to/local/model``
          - ``relative/path/to/local/model``
          - ``s3://my_bucket/path/to/model``
          - ``runs:/<mlflow_run_id>/run-relative/path/to/model``
          - ``models:/<model_name>/<model_version>``
          - ``models:/<model_name>/<stage>``
  
          For more information about supported URI schemes, see `Referencing Artifacts \
          <https://www.mlflow.org/docs/latest/concepts.html#artifact-locations>`_.
      :param dst_path: The local filesystem path to which to download the model artifact.
          This directory must already exist. If unspecified, a local output path will be created.
      :param kwargs: kwargs to pass to ``torch.load`` method.
  
      :return: A PyTorch model.
  
      .. code-block:: python
          :caption: Example
  
          import torch
          import mlflow.pytorch
  
  
          model = nn.Linear(1, 1)
  
          # Log the model
          with mlflow.start_run() as run:
              mlflow.pytorch.log_model(model, "model")
  
          # Inference after loading the logged model
          model_uri = f"runs:/{run.info.run_id}/model"
          loaded_model = mlflow.pytorch.load_model(model_uri)
          for x in [4.0, 6.0, 30.0]:
              X = torch.Tensor([[x]])
              y_pred = loaded_model(X)
              print(f"predict X: {x}, y_pred: {y_pred.data.item():.2f}")
  
      .. code-block:: text
          :caption: Output
  
          predict X: 4.0, y_pred: 7.57
          predict X: 6.0, y_pred: 11.64
          predict X: 30.0, y_pred: 60.48
      """
      import torch
  
      local_model_path = _download_artifact_from_uri(
          artifact_uri=model_uri, output_path=dst_path
      )
      pytorch_conf = _get_flavor_configuration(
          model_path=local_model_path, flavor_name=FLAVOR_NAME
      )
      _add_code_from_conf_to_system_path(local_model_path, pytorch_conf)
  
      if torch.__version__ != pytorch_conf["pytorch_version"]:
          _logger.warning(
              "Stored model version '%s' does not match installed PyTorch version '%s'",
              pytorch_conf["pytorch_version"],
              torch.__version__,
          )
      torch_model_artifacts_path = os.path.join(
          local_model_path, pytorch_conf["model_data"]
      )
      return _load_model(path=torch_model_artifacts_path, **kwargs)
  
imports: import atexit
  import importlib
  import logging
  import os
  import posixpath
  import shutil
  import warnings
  from functools import partial
  from typing import Any, Dict, Optional
  import numpy
  import pandas
  import yaml
  from packaging.version import Version
  import mlflow
  from mlflow import pyfunc
  from mlflow.environment_variables import MLFLOW_DEFAULT_PREDICTION_DEVICE
  from mlflow.exceptions import MlflowException
  from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS
  from mlflow.models import Model, ModelSignature
  from mlflow.models.model import MLMODEL_FILE_NAME
  from mlflow.models.signature import _infer_signature_from_input_example
  from mlflow.models.utils import ModelInputExample, _save_example
  from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
  from mlflow.pytorch import pickle_module
  from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
  from mlflow.tracking.artifact_utils import _download_artifact_from_uri
  from mlflow.utils.autologging_utils import autologging_integration, safe_patch
  from mlflow.utils.docstring_utils import LOG_MODEL_PARAM_DOCS, format_docstring
  from mlflow.utils.environment import _CONDA_ENV_FILE_NAME, _CONSTRAINTS_FILE_NAME, _PYTHON_ENV_FILE_NAME, _REQUIREMENTS_FILE_NAME, _mlflow_conda_env, _process_conda_env, _process_pip_requirements, _PythonEnv, _validate_env_arguments
  from mlflow.utils.file_utils import TempDir, get_total_file_size, write_to
  from mlflow.utils.model_utils import _add_code_from_conf_to_system_path, _get_flavor_configuration, _validate_and_copy_code_paths, _validate_and_prepare_target_save_path
  from mlflow.utils.requirements_utils import _get_pinned_requirement
label_desc: Use image node_id 6 for calling a global function with example usage: load_model(model_uri, dst_path) and returns: _load_model