code: def log_model(
      pytorch_model,
      artifact_path,
      conda_env=None,
      code_paths=None,
      pickle_module=None,
      registered_model_name=None,
      signature: ModelSignature = None,
      input_example: ModelInputExample = None,
      await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
      requirements_file=None,
      extra_files=None,
      pip_requirements=None,
      extra_pip_requirements=None,
      metadata=None,
      **kwargs,
  ):
      """
      Log a PyTorch model as an MLflow artifact for the current run.
  
          .. warning::
  
              Log the model with a signature to avoid inference errors.
              If the model is logged without a signature, the MLflow Model Server relies on the
              default inferred data type from NumPy. However, PyTorch often expects different
              defaults, particularly when parsing floats. You must include the signature to ensure
              that the model is logged with the correct data type so that the MLflow model server
              can correctly provide valid input.
  
      :param pytorch_model: PyTorch model to be saved. Can be either an eager model (subclass of
          ``torch.nn.Module``) or scripted model prepared via ``torch.jit.script`` or
          ``torch.jit.trace``.
  
          The model accept a single ``torch.FloatTensor`` as input and produce a single output tensor.
  
          If saving 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 artifact_path: Run-relative artifact path.
      :param conda_env: {{ conda_env }}
      :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 PyTorch model at load
          time.
      :param registered_model_name: If given, create a model version under ``registered_model_name``,
          also create a registered model if one with the given name does not exist.
      :param signature: {{ signature }}
      :param input_example: {{ input_example }}
      :param await_registration_for: Number of seconds to wait for the model version to finish
          being created and is in ``READY`` status. By default, the function waits for five minutes.
          Specify 0 or None to skip waiting.
  
      :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, if ``None``, no
          extra files are added to the model. Remote URIs are resolved to absolute filesystem paths.
          For example, consider the following ``extra_files`` list:
  
          .. code-block:: python
  
              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.
  
      :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.
      :return: A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
          metadata of the logged model.
  
      .. code-block:: python
          :caption: Example
  
          import numpy as np
          import torch
          import mlflow
          from mlflow import MlflowClient
          from mlflow.models import infer_signature
  
          # Define model, loss, and optimizer
          model = nn.Linear(1, 1)
          criterion = torch.nn.MSELoss()
          optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
  
          # Create training data with relationship y = 2X
          X = torch.arange(1.0, 26.0).reshape(-1, 1)
          y = X * 2
  
          # Training loop
          epochs = 250
          for epoch in range(epochs):
              # Forward pass: Compute predicted y by passing X to the model
              y_pred = model(X)
  
              # Compute the loss
              loss = criterion(y_pred, y)
  
              # Zero gradients, perform a backward pass, and update the weights.
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
  
          # Create model signature
          signature = infer_signature(X.numpy(), model(X).detach().numpy())
  
          # Log the model
          with mlflow.start_run() as run:
              mlflow.pytorch.log_model(model, "model")
  
              # convert to scripted model and log the model
              scripted_pytorch_model = torch.jit.script(model)
              mlflow.pytorch.log_model(scripted_pytorch_model, "scripted_model")
  
          # Fetch the logged model artifacts
          print(f"run_id: {run.info.run_id}")
          for artifact_path in ["model/data", "scripted_model/data"]:
              artifacts = [
                  f.path for f in MlflowClient().list_artifacts(run.info.run_id, artifact_path)
              ]
              print(f"artifacts: {artifacts}")
  
      .. code-block:: text
          :caption: Output
  
          run_id: 1a1ec9e413ce48e9abf9aec20efd6f71
          artifacts: ['model/data/model.pth',
                      'model/data/pickle_module_info.txt']
          artifacts: ['scripted_model/data/model.pth',
                      'scripted_model/data/pickle_module_info.txt']
  
      .. figure:: ../_static/images/pytorch_logged_models.png
  
          PyTorch logged models
      """
      pickle_module = pickle_module or mlflow_pytorch_pickle_module
      return Model.log(
          artifact_path=artifact_path,
          flavor=mlflow.pytorch,
          pytorch_model=pytorch_model,
          conda_env=conda_env,
          code_paths=code_paths,
          pickle_module=pickle_module,
          registered_model_name=registered_model_name,
          signature=signature,
          input_example=input_example,
          await_registration_for=await_registration_for,
          requirements_file=requirements_file,
          extra_files=extra_files,
          pip_requirements=pip_requirements,
          extra_pip_requirements=extra_pip_requirements,
          metadata=metadata,
          **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 3 for calling a global function with example usage: log_model(pytorch_model, artifact_path, conda_env, code_paths, pickle_module, registered_model_name, signature, input_example, await_registration_for, requirements_file, extra_files, pip_requirements, extra_pip_requirements, metadata) and returns: Model