code: def save_model(
      pytorch_model,
      path,
      conda_env=None,
      mlflow_model=None,
      code_paths=None,
      pickle_module=None,
      signature: ModelSignature = None,
      input_example: ModelInputExample = None,
      requirements_file=None,
      extra_files=None,
      pip_requirements=None,
      extra_pip_requirements=None,
      metadata=None,
      **kwargs,
  ):
      """
      Save a PyTorch model to a path on the local file system.
  
      :param pytorch_model: PyTorch model to be saved. Can be either an eager model (subclass of
          ``torch.nn.Module``) or a scripted model prepared via ``torch.jit.script`` or
          ``torch.jit.trace``.
  
          To save an eager model, any code dependencies of the model's class, including the class
          definition itself, should be included in one of the following locations:
  
              - The package(s) listed in the model's Conda environment, specified by the ``conda_env``
                  parameter.
              - One or more of the files specified by the ``code_paths`` parameter.
  
      :param path: Local path where the model is to be saved.
      :param conda_env: {{ conda_env }}
      :param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
      :param code_paths: A list of local filesystem paths to Python file dependencies (or directories
          containing file dependencies). These files are *prepended* to the system path when the model
          is loaded.
      :param pickle_module: The module that PyTorch should use to serialize ("pickle") the specified
          ``pytorch_model``. This is passed as the ``pickle_module`` parameter to ``torch.save()``. By
          default, this module is also used to deserialize ("unpickle") the model at loading time.
      :param signature: {{ signature }}
      :param input_example: {{ input_example }}
      :param requirements_file:
  
          .. warning::
  
              ``requirements_file`` has been deprecated. Please use ``pip_requirements`` instead.
  
          A string containing the path to requirements file. Remote URIs are resolved to absolute
          filesystem paths. For example, consider the following ``requirements_file`` string:
  
          .. code-block:: python
  
              requirements_file = "s3://my-bucket/path/to/my_file"
  
          In this case, the ``"my_file"`` requirements file is downloaded from S3. If ``None``,
          no requirements file is added to the model.
  
      :param extra_files: A list containing the paths to corresponding extra files. Remote URIs
          are resolved to absolute filesystem paths.
          For example, consider the following ``extra_files`` list -
  
          extra_files = ["s3://my-bucket/path/to/my_file1",
                      "s3://my-bucket/path/to/my_file2"]
  
          In this case, the ``"my_file1 & my_file2"`` extra file is downloaded from S3.
  
          If ``None``, no extra files are added to the model.
      :param pip_requirements: {{ pip_requirements }}
      :param extra_pip_requirements: {{ extra_pip_requirements }}
      :param metadata: Custom metadata dictionary passed to the model and stored in the MLmodel file.
  
          .. Note:: Experimental: This parameter may change or be removed in a future
                                  release without warning.
      :param kwargs: kwargs to pass to ``torch.save`` method.
  
      .. code-block:: python
          :caption: Example
  
          import os
          import mlflow
          import torch
  
  
          model = nn.Linear(1, 1)
  
          # Save PyTorch models to current working directory
          with mlflow.start_run() as run:
              mlflow.pytorch.save_model(model, "model")
  
              # Convert to a scripted model and save it
              scripted_pytorch_model = torch.jit.script(model)
              mlflow.pytorch.save_model(scripted_pytorch_model, "scripted_model")
  
          # Load each saved model for inference
          for model_path in ["model", "scripted_model"]:
              model_uri = f"{os.getcwd()}/{model_path}"
              loaded_model = mlflow.pytorch.load_model(model_uri)
              print(f"Loaded {model_path}:")
              for x in [6.0, 8.0, 12.0, 30.0]:
                  X = torch.Tensor([[x]])
                  y_pred = loaded_model(X)
                  print(f"predict X: {x}, y_pred: {y_pred.data.item():.2f}")
              print("--")
  
      .. code-block:: text
          :caption: Output
  
          Loaded model:
          predict X: 6.0, y_pred: 11.90
          predict X: 8.0, y_pred: 15.92
          predict X: 12.0, y_pred: 23.96
          predict X: 30.0, y_pred: 60.13
          --
          Loaded scripted_model:
          predict X: 6.0, y_pred: 11.90
          predict X: 8.0, y_pred: 15.92
          predict X: 12.0, y_pred: 23.96
          predict X: 30.0, y_pred: 60.13
  
      """
      import torch
  
      _validate_env_arguments(
          conda_env, pip_requirements, extra_pip_requirements
      )
  
      pickle_module = pickle_module or mlflow_pytorch_pickle_module
  
      if not isinstance(pytorch_model, torch.nn.Module):
          raise TypeError(
              "Argument 'pytorch_model' should be a torch.nn.Module"
          )
      path = os.path.abspath(path)
      _validate_and_prepare_target_save_path(path)
  
      if signature is None and input_example is not None:
          wrapped_model = _PyTorchWrapper(pytorch_model, device="cpu")
          signature = _infer_signature_from_input_example(
              input_example, wrapped_model
          )
      elif signature is False:
          signature = None
  
      if mlflow_model is None:
          mlflow_model = Model()
      if signature is not None:
          mlflow_model.signature = signature
      if input_example is not None:
          _save_example(mlflow_model, input_example, path)
      if metadata is not None:
          mlflow_model.metadata = metadata
  
      code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
  
      model_data_subpath = "data"
      model_data_path = os.path.join(path, model_data_subpath)
      os.makedirs(model_data_path)
  
      # Persist the pickle module name as a file in the model's `data` directory. This is necessary
      # because the `data` directory is the only available parameter to `_load_pyfunc`, and it
      # does not contain the MLmodel configuration; therefore, it is not sufficient to place
      # the module name in the MLmodel
      #
      # TODO: Stop persisting this information to the filesystem once we have a mechanism for
      # supplying the MLmodel configuration to `mlflow.pytorch._load_pyfunc`
      pickle_module_path = os.path.join(
          model_data_path, _PICKLE_MODULE_INFO_FILE_NAME
      )
      with open(pickle_module_path, "w") as f:
          f.write(pickle_module.__name__)
      # Save pytorch model
      model_path = os.path.join(
          model_data_path, _SERIALIZED_TORCH_MODEL_FILE_NAME
      )
      if isinstance(pytorch_model, torch.jit.ScriptModule):
          torch.jit.ScriptModule.save(pytorch_model, model_path)
      else:
          torch.save(
              pytorch_model,
              model_path,
              pickle_module=pickle_module,
              **kwargs,
          )
  
      torchserve_artifacts_config = {}
  
      if extra_files:
          torchserve_artifacts_config[_EXTRA_FILES_KEY] = []
          if not isinstance(extra_files, list):
              raise TypeError("Extra files argument should be a list")
  
          with TempDir() as tmp_extra_files_dir:
              for extra_file in extra_files:
                  _download_artifact_from_uri(
                      artifact_uri=extra_file,
                      output_path=tmp_extra_files_dir.path(),
                  )
                  rel_path = posixpath.join(
                      _EXTRA_FILES_KEY, os.path.basename(extra_file)
                  )
                  torchserve_artifacts_config[_EXTRA_FILES_KEY].append(
                      {"path": rel_path}
                  )
              shutil.move(
                  tmp_extra_files_dir.path(),
                  posixpath.join(path, _EXTRA_FILES_KEY),
              )
  
      if requirements_file:
          warnings.warn(
              "`requirements_file` has been deprecated. Please use `pip_requirements` instead.",
              FutureWarning,
              stacklevel=2,
          )
  
          if not isinstance(requirements_file, str):
              raise TypeError(
                  "Path to requirements file should be a string"
              )
  
          with TempDir() as tmp_requirements_dir:
              _download_artifact_from_uri(
                  artifact_uri=requirements_file,
                  output_path=tmp_requirements_dir.path(),
              )
              rel_path = os.path.basename(requirements_file)
              torchserve_artifacts_config[_REQUIREMENTS_FILE_KEY] = {
                  "path": rel_path
              }
              shutil.move(tmp_requirements_dir.path(rel_path), path)
  
      mlflow_model.add_flavor(
          FLAVOR_NAME,
          model_data=model_data_subpath,
          pytorch_version=str(torch.__version__),
          code=code_dir_subpath,
          **torchserve_artifacts_config,
      )
      pyfunc.add_to_model(
          mlflow_model,
          loader_module="mlflow.pytorch",
          data=model_data_subpath,
          pickle_module_name=pickle_module.__name__,
          code=code_dir_subpath,
          conda_env=_CONDA_ENV_FILE_NAME,
          python_env=_PYTHON_ENV_FILE_NAME,
          model_config={"device": None},
      )
      if size := get_total_file_size(path):
          mlflow_model.model_size_bytes = size
      mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
  
      if conda_env is None:
          if pip_requirements is None:
              default_reqs = get_default_pip_requirements()
              # To ensure `_load_pyfunc` can successfully load the model during the dependency
              # inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
              inferred_reqs = mlflow.models.infer_pip_requirements(
                  model_data_path,
                  FLAVOR_NAME,
                  fallback=default_reqs,
              )
              default_reqs = sorted(
                  set(inferred_reqs).union(default_reqs)
              )
          else:
              default_reqs = None
          (
              conda_env,
              pip_requirements,
              pip_constraints,
          ) = _process_pip_requirements(
              default_reqs,
              pip_requirements,
              extra_pip_requirements,
          )
      else:
          (
              conda_env,
              pip_requirements,
              pip_constraints,
          ) = _process_conda_env(conda_env)
  
      with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
          yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
  
      # Save `constraints.txt` if necessary
      if pip_constraints:
          write_to(
              os.path.join(path, _CONSTRAINTS_FILE_NAME),
              "\n".join(pip_constraints),
          )
  
      if not requirements_file:
          # Save `requirements.txt`
          write_to(
              os.path.join(path, _REQUIREMENTS_FILE_NAME),
              "\n".join(pip_requirements),
          )
  
      _PythonEnv.current().to_yaml(
          os.path.join(path, _PYTHON_ENV_FILE_NAME)
      )
  
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 4 for calling a global function with example usage: save_model(pytorch_model, path, conda_env, mlflow_model, code_paths, pickle_module, signature, input_example, requirements_file, extra_files, pip_requirements, extra_pip_requirements, metadata) without return types