ultralytics 8.1.8 new model.save('filename.pt') method (#7886)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

docs/overrides/partials/comments.html

@@ -9,7 +9,8 @@
         data-emit-metadata="0"
         data-input-position="top"
         data-lang="en"
-        data-mapping="title"
+        data-loading="lazy"
+        data-mapping="pathname"
         data-reactions-enabled="1"
         data-repo="ultralytics/ultralytics"
         data-repo-id="R_kgDOH-jzvQ"

mkdocs.yml

@@ -179,7 +179,7 @@ nav:
       - NEW 🚀 Explorer:
           - datasets/explorer/index.md
       - Languages:
-          - 🇬🇧&nbsp English: https://docs.ultralytics.com/
+          - 🇬🇧&nbsp English: https://ultralytics.com/docs/
           - 🇨🇳&nbsp 简体中文: https://docs.ultralytics.com/zh/
           - 🇰🇷&nbsp 한국어: https://docs.ultralytics.com/ko/
           - 🇯🇵&nbsp 日本語: https://docs.ultralytics.com/ja/

ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.1.7"
+__version__ = "8.1.8"
 
 from ultralytics.data.explorer.explorer import Explorer
 from ultralytics.models import RTDETR, SAM, YOLO

ultralytics/engine/model.py

@@ -6,60 +6,98 @@ from pathlib import Path
 from typing import Union
 
 from ultralytics.cfg import TASK2DATA, get_cfg, get_save_dir
+from ultralytics.hub.utils import HUB_WEB_ROOT
 from ultralytics.nn.tasks import attempt_load_one_weight, guess_model_task, nn, yaml_model_load
 from ultralytics.utils import ASSETS, DEFAULT_CFG_DICT, LOGGER, RANK, SETTINGS, callbacks, checks, emojis, yaml_load
-from ultralytics.hub.utils import HUB_WEB_ROOT
 
 
 class Model(nn.Module):
     """
-    A base class to unify APIs for all models.
+    A base class for implementing YOLO models, unifying APIs across different model types.
+
+    This class provides a common interface for various operations related to YOLO models, such as training,
+    validation, prediction, exporting, and benchmarking. It handles different types of models, including those
+    loaded from local files, Ultralytics HUB, or Triton Server. The class is designed to be flexible and
+    extendable for different tasks and model configurations.
 
     Args:
-        model (str, Path): Path to the model file to load or create.
+        model (Union[str, Path], optional): Path or name of the model to load or create. This can be a local file
-        task (Any, optional): Task type for the YOLO model. Defaults to None.
+            path, a model name from Ultralytics HUB, or a Triton Server model. Defaults to 'yolov8n.pt'.
+        task (Any, optional): The task type associated with the YOLO model. This can be used to specify the model's
+            application domain, such as object detection, segmentation, etc. Defaults to None.
+        verbose (bool, optional): If True, enables verbose output during the model's operations. Defaults to False.
 
     Attributes:
-        predictor (Any): The predictor object.
+        callbacks (dict): A dictionary of callback functions for various events during model operations.
-        model (Any): The model object.
+        predictor (BasePredictor): The predictor object used for making predictions.
-        trainer (Any): The trainer object.
+        model (nn.Module): The underlying PyTorch model.
-        task (str): The type of model task.
+        trainer (BaseTrainer): The trainer object used for training the model.
-        ckpt (Any): The checkpoint object if the model loaded from *.pt file.
+        ckpt (dict): The checkpoint data if the model is loaded from a *.pt file.
-        cfg (str): The model configuration if loaded from *.yaml file.
+        cfg (str): The configuration of the model if loaded from a *.yaml file.
-        ckpt_path (str): The checkpoint file path.
+        ckpt_path (str): The path to the checkpoint file.
-        overrides (dict): Overrides for the trainer object.
+        overrides (dict): A dictionary of overrides for model configuration.
-        metrics (Any): The data for metrics.
+        metrics (dict): The latest training/validation metrics.
+        session (HUBTrainingSession): The Ultralytics HUB session, if applicable.
+        task (str): The type of task the model is intended for.
+        model_name (str): The name of the model.
 
     Methods:
-        __call__(source=None, stream=False, **kwargs):
+        __call__: Alias for the predict method, enabling the model instance to be callable.
-            Alias for the predict method.
+        _new: Initializes a new model based on a configuration file.
-        _new(cfg:str, verbose:bool=True) -> None:
+        _load: Loads a model from a checkpoint file.
-            Initializes a new model and infers the task type from the model definitions.
+        _check_is_pytorch_model: Ensures that the model is a PyTorch model.
-        _load(weights:str, task:str='') -> None:
+        reset_weights: Resets the model's weights to their initial state.
-            Initializes a new model and infers the task type from the model head.
+        load: Loads model weights from a specified file.
-        _check_is_pytorch_model() -> None:
+        save: Saves the current state of the model to a file.
-            Raises TypeError if the model is not a PyTorch model.
+        info: Logs or returns information about the model.
-        reset() -> None:
+        fuse: Fuses Conv2d and BatchNorm2d layers for optimized inference.
-            Resets the model modules.
+        predict: Performs object detection predictions.
-        info(verbose:bool=False) -> None:
+        track: Performs object tracking.
-            Logs the model info.
+        val: Validates the model on a dataset.
-        fuse() -> None:
+        benchmark: Benchmarks the model on various export formats.
-            Fuses the model for faster inference.
+        export: Exports the model to different formats.
-        predict(source=None, stream=False, **kwargs) -> List[ultralytics.engine.results.Results]:
+        train: Trains the model on a dataset.
-            Performs prediction using the YOLO model.
+        tune: Performs hyperparameter tuning.
+        _apply: Applies a function to the model's tensors.
+        add_callback: Adds a callback function for an event.
+        clear_callback: Clears all callbacks for an event.
+        reset_callbacks: Resets all callbacks to their default functions.
+        _get_hub_session: Retrieves or creates an Ultralytics HUB session.
+        is_triton_model: Checks if a model is a Triton Server model.
+        is_hub_model: Checks if a model is an Ultralytics HUB model.
+        _reset_ckpt_args: Resets checkpoint arguments when loading a PyTorch model.
+        _smart_load: Loads the appropriate module based on the model task.
+        task_map: Provides a mapping from model tasks to corresponding classes.
 
-    Returns:
+    Raises:
-        list(ultralytics.engine.results.Results): The prediction results.
+        FileNotFoundError: If the specified model file does not exist or is inaccessible.
+        ValueError: If the model file or configuration is invalid or unsupported.
+        ImportError: If required dependencies for specific model types (like HUB SDK) are not installed.
+        TypeError: If the model is not a PyTorch model when required.
+        AttributeError: If required attributes or methods are not implemented or available.
+        NotImplementedError: If a specific model task or mode is not supported.
     """
 
     def __init__(self, model: Union[str, Path] = "yolov8n.pt", task=None, verbose=False) -> None:
         """
-        Initializes the YOLO model.
+        Initializes a new instance of the YOLO model class.
+
+        This constructor sets up the model based on the provided model path or name. It handles various types of model
+        sources, including local files, Ultralytics HUB models, and Triton Server models. The method initializes several
+        important attributes of the model and prepares it for operations like training, prediction, or export.
 
         Args:
-            model (Union[str, Path], optional): Path or name of the model to load or create. Defaults to 'yolov8n.pt'.
+            model (Union[str, Path], optional): The path or model file to load or create. This can be a local
-            task (Any, optional): Task type for the YOLO model. Defaults to None.
+                file path, a model name from Ultralytics HUB, or a Triton Server model. Defaults to 'yolov8n.pt'.
-            verbose (bool, optional): Whether to enable verbose mode.
+            task (Any, optional): The task type associated with the YOLO model, specifying its application domain.
+                Defaults to None.
+            verbose (bool, optional): If True, enables verbose output during the model's initialization and subsequent
+                operations. Defaults to False.
+
+        Raises:
+            FileNotFoundError: If the specified model file does not exist or is inaccessible.
+            ValueError: If the model file or configuration is invalid or unsupported.
+            ImportError: If required dependencies for specific model types (like HUB SDK) are not installed.
         """
         super().__init__()
         self.callbacks = callbacks.get_default_callbacks()
@@ -98,7 +136,22 @@ class Model(nn.Module):
         self.model_name = model
 
     def __call__(self, source=None, stream=False, **kwargs):
-        """Calls the predict() method with given arguments to perform object detection."""
+        """
+        An alias for the predict method, enabling the model instance to be callable.
+
+        This method simplifies the process of making predictions by allowing the model instance to be called directly
+        with the required arguments for prediction.
+
+        Args:
+            source (str | int | PIL.Image | np.ndarray, optional): The source of the image for making predictions.
+                Accepts various types, including file paths, URLs, PIL images, and numpy arrays. Defaults to None.
+            stream (bool, optional): If True, treats the input source as a continuous stream for predictions.
+                Defaults to False.
+            **kwargs (dict): Additional keyword arguments for configuring the prediction process.
+
+        Returns:
+            (List[ultralytics.engine.results.Results]): A list of prediction results, encapsulated in the Results class.
+        """
         return self.predict(source, stream, **kwargs)
 
     @staticmethod
@@ -185,7 +238,19 @@ class Model(nn.Module):
             )
 
     def reset_weights(self):
-        """Resets the model modules parameters to randomly initialized values, losing all training information."""
+        """
+        Resets the model parameters to randomly initialized values, effectively discarding all training information.
+
+        This method iterates through all modules in the model and resets their parameters if they have a
+        'reset_parameters' method. It also ensures that all parameters have 'requires_grad' set to True, enabling them
+        to be updated during training.
+
+        Returns:
+            self (ultralytics.engine.model.Model): The instance of the class with reset weights.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
+        """
         self._check_is_pytorch_model()
         for m in self.model.modules():
             if hasattr(m, "reset_parameters"):
@@ -195,42 +260,94 @@ class Model(nn.Module):
         return self
 
     def load(self, weights="yolov8n.pt"):
-        """Transfers parameters with matching names and shapes from 'weights' to model."""
+        """
+        Loads parameters from the specified weights file into the model.
+
+        This method supports loading weights from a file or directly from a weights object. It matches parameters by
+        name and shape and transfers them to the model.
+
+        Args:
+            weights (str | Path): Path to the weights file or a weights object. Defaults to 'yolov8n.pt'.
+
+        Returns:
+            self (ultralytics.engine.model.Model): The instance of the class with loaded weights.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
+        """
         self._check_is_pytorch_model()
         if isinstance(weights, (str, Path)):
             weights, self.ckpt = attempt_load_one_weight(weights)
         self.model.load(weights)
         return self
 
-    def info(self, detailed=False, verbose=True):
+    def save(self, filename="model.pt"):
         """
-        Logs model info.
+        Saves the current model state to a file.
+
+        This method exports the model's checkpoint (ckpt) to the specified filename.
 
         Args:
-            detailed (bool): Show detailed information about model.
+            filename (str): The name of the file to save the model to. Defaults to 'model.pt'.
-            verbose (bool): Controls verbosity.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
+        """
+        self._check_is_pytorch_model()
+        import torch
+
+        torch.save(self.ckpt, filename)
+
+    def info(self, detailed=False, verbose=True):
+        """
+        Logs or returns model information.
+
+        This method provides an overview or detailed information about the model, depending on the arguments passed.
+        It can control the verbosity of the output.
+
+        Args:
+            detailed (bool): If True, shows detailed information about the model. Defaults to False.
+            verbose (bool): If True, prints the information. If False, returns the information. Defaults to True.
+
+        Returns:
+            (list): Various types of information about the model, depending on the 'detailed' and 'verbose' parameters.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
         """
         self._check_is_pytorch_model()
         return self.model.info(detailed=detailed, verbose=verbose)
 
     def fuse(self):
-        """Fuse PyTorch Conv2d and BatchNorm2d layers."""
+        """
+        Fuses Conv2d and BatchNorm2d layers in the model.
+
+        This method optimizes the model by fusing Conv2d and BatchNorm2d layers, which can improve inference speed.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
+        """
         self._check_is_pytorch_model()
         self.model.fuse()
 
     def embed(self, source=None, stream=False, **kwargs):
         """
-        Calls the predict() method and returns image embeddings.
+        Generates image embeddings based on the provided source.
+
+        This method is a wrapper around the 'predict()' method, focusing on generating embeddings from an image source.
+        It allows customization of the embedding process through various keyword arguments.
 
         Args:
-            source (str | int | PIL | np.ndarray): The source of the image to make predictions on.
+            source (str | int | PIL.Image | np.ndarray): The source of the image for generating embeddings.
-                Accepts all source types accepted by the YOLO model.
+                The source can be a file path, URL, PIL image, numpy array, etc. Defaults to None.
-            stream (bool): Whether to stream the predictions or not. Defaults to False.
+            stream (bool): If True, predictions are streamed. Defaults to False.
-            **kwargs : Additional keyword arguments passed to the predictor.
+            **kwargs (dict): Additional keyword arguments for configuring the embedding process.
-                Check the 'configuration' section in the documentation for all available options.
 
         Returns:
-            (List[torch.Tensor]): A list of image embeddings.
+            (List[torch.Tensor]): A list containing the image embeddings.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
         """
         if not kwargs.get("embed"):
             kwargs["embed"] = [len(self.model.model) - 2]  # embed second-to-last layer if no indices passed
@@ -238,18 +355,32 @@ class Model(nn.Module):
 
     def predict(self, source=None, stream=False, predictor=None, **kwargs):
         """
-        Perform prediction using the YOLO model.
+        Performs predictions on the given image source using the YOLO model.
+
+        This method facilitates the prediction process, allowing various configurations through keyword arguments.
+        It supports predictions with custom predictors or the default predictor method. The method handles different
+        types of image sources and can operate in a streaming mode. It also provides support for SAM-type models
+        through 'prompts'.
+
+        The method sets up a new predictor if not already present and updates its arguments with each call.
+        It also issues a warning and uses default assets if the 'source' is not provided. The method determines if it
+        is being called from the command line interface and adjusts its behavior accordingly, including setting defaults
+        for confidence threshold and saving behavior.
 
         Args:
-            source (str | int | PIL | np.ndarray): The source of the image to make predictions on.
+            source (str | int | PIL.Image | np.ndarray, optional): The source of the image for making predictions.
-                Accepts all source types accepted by the YOLO model.
+                Accepts various types, including file paths, URLs, PIL images, and numpy arrays. Defaults to ASSETS.
-            stream (bool): Whether to stream the predictions or not. Defaults to False.
+            stream (bool, optional): Treats the input source as a continuous stream for predictions. Defaults to False.
-            predictor (BasePredictor): Customized predictor.
+            predictor (BasePredictor, optional): An instance of a custom predictor class for making predictions.
-            **kwargs : Additional keyword arguments passed to the predictor.
+                If None, the method uses a default predictor. Defaults to None.
-                Check the 'configuration' section in the documentation for all available options.
+            **kwargs (dict): Additional keyword arguments for configuring the prediction process. These arguments allow
+                for further customization of the prediction behavior.
 
         Returns:
-            (List[ultralytics.engine.results.Results]): The prediction results.
+            (List[ultralytics.engine.results.Results]): A list of prediction results, encapsulated in the Results class.
+
+        Raises:
+            AttributeError: If the predictor is not properly set up.
         """
         if source is None:
             source = ASSETS
@@ -276,16 +407,28 @@ class Model(nn.Module):
 
     def track(self, source=None, stream=False, persist=False, **kwargs):
         """
-        Perform object tracking on the input source using the registered trackers.
+        Conducts object tracking on the specified input source using the registered trackers.
+
+        This method performs object tracking using the model's predictors and optionally registered trackers. It is
+        capable of handling different types of input sources such as file paths or video streams. The method supports
+        customization of the tracking process through various keyword arguments. It registers trackers if they are not
+        already present and optionally persists them based on the 'persist' flag.
+
+        The method sets a default confidence threshold specifically for ByteTrack-based tracking, which requires low
+        confidence predictions as input. The tracking mode is explicitly set in the keyword arguments.
 
         Args:
-            source (str, optional): The input source for object tracking. Can be a file path or a video stream.
+            source (str, optional): The input source for object tracking. It can be a file path, URL, or video stream.
-            stream (bool, optional): Whether the input source is a video stream. Defaults to False.
+            stream (bool, optional): Treats the input source as a continuous video stream. Defaults to False.
-            persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False.
+            persist (bool, optional): Persists the trackers between different calls to this method. Defaults to False.
-            **kwargs (optional): Additional keyword arguments for the tracking process.
+            **kwargs (dict): Additional keyword arguments for configuring the tracking process. These arguments allow
+                for further customization of the tracking behavior.
 
         Returns:
-            (List[ultralytics.engine.results.Results]): The tracking results.
+            (List[ultralytics.engine.results.Results]): A list of tracking results, encapsulated in the Results class.
+
+        Raises:
+            AttributeError: If the predictor does not have registered trackers.
         """
         if not hasattr(self.predictor, "trackers"):
             from ultralytics.trackers import register_tracker
@@ -297,11 +440,28 @@ class Model(nn.Module):
 
     def val(self, validator=None, **kwargs):
         """
-        Validate a model on a given dataset.
+        Validates the model using a specified dataset and validation configuration.
+
+        This method facilitates the model validation process, allowing for a range of customization through various
+        settings and configurations. It supports validation with a custom validator or the default validation approach.
+        The method combines default configurations, method-specific defaults, and user-provided arguments to configure
+        the validation process. After validation, it updates the model's metrics with the results obtained from the
+        validator.
+
+        The method supports various arguments that allow customization of the validation process. For a comprehensive
+        list of all configurable options, users should refer to the 'configuration' section in the documentation.
 
         Args:
-            validator (BaseValidator): Customized validator.
+            validator (BaseValidator, optional): An instance of a custom validator class for validating the model. If
-            **kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs
+                None, the method uses a default validator. Defaults to None.
+            **kwargs (dict): Arbitrary keyword arguments representing the validation configuration. These arguments are
+                used to customize various aspects of the validation process.
+
+        Returns:
+            (dict): Validation metrics obtained from the validation process.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
         """
         custom = {"rect": True}  # method defaults
         args = {**self.overrides, **custom, **kwargs, "mode": "val"}  # highest priority args on the right
@@ -313,10 +473,26 @@ class Model(nn.Module):
 
     def benchmark(self, **kwargs):
         """
-        Benchmark a model on all export formats.
+        Benchmarks the model across various export formats to evaluate performance.
+
+        This method assesses the model's performance in different export formats, such as ONNX, TorchScript, etc.
+        It uses the 'benchmark' function from the ultralytics.utils.benchmarks module. The benchmarking is configured
+        using a combination of default configuration values, model-specific arguments, method-specific defaults, and
+        any additional user-provided keyword arguments.
+
+        The method supports various arguments that allow customization of the benchmarking process, such as dataset
+        choice, image size, precision modes, device selection, and verbosity. For a comprehensive list of all
+        configurable options, users should refer to the 'configuration' section in the documentation.
 
         Args:
-            **kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs
+            **kwargs (dict): Arbitrary keyword arguments to customize the benchmarking process. These are combined with
+                default configurations, model-specific arguments, and method defaults.
+
+        Returns:
+            (dict): A dictionary containing the results of the benchmarking process.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
         """
         self._check_is_pytorch_model()
         from ultralytics.utils.benchmarks import benchmark
@@ -335,10 +511,24 @@ class Model(nn.Module):
 
     def export(self, **kwargs):
         """
-        Export model.
+        Exports the model to a different format suitable for deployment.
+
+        This method facilitates the export of the model to various formats (e.g., ONNX, TorchScript) for deployment
+        purposes. It uses the 'Exporter' class for the export process, combining model-specific overrides, method
+        defaults, and any additional arguments provided. The combined arguments are used to configure export settings.
+
+        The method supports a wide range of arguments to customize the export process. For a comprehensive list of all
+        possible arguments, refer to the 'configuration' section in the documentation.
 
         Args:
-            **kwargs : Any other args accepted by the Exporter. To see all args check 'configuration' section in docs.
+            **kwargs (dict): Arbitrary keyword arguments to customize the export process. These are combined with the
+                model's overrides and method defaults.
+
+        Returns:
+            (object): The exported model in the specified format, or an object related to the export process.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
         """
         self._check_is_pytorch_model()
         from .exporter import Exporter
@@ -349,11 +539,31 @@ class Model(nn.Module):
 
     def train(self, trainer=None, **kwargs):
         """
-        Trains the model on a given dataset.
+        Trains the model using the specified dataset and training configuration.
+
+        This method facilitates model training with a range of customizable settings and configurations. It supports
+        training with a custom trainer or the default training approach defined in the method. The method handles
+        different scenarios, such as resuming training from a checkpoint, integrating with Ultralytics HUB, and
+        updating model and configuration after training.
+
+        When using Ultralytics HUB, if the session already has a loaded model, the method prioritizes HUB training
+        arguments and issues a warning if local arguments are provided. It checks for pip updates and combines default
+        configurations, method-specific defaults, and user-provided arguments to configure the training process. After
+        training, it updates the model and its configurations, and optionally attaches metrics.
 
         Args:
-            trainer (BaseTrainer, optional): Customized trainer.
+            trainer (BaseTrainer, optional): An instance of a custom trainer class for training the model. If None, the
-            **kwargs (Any): Any number of arguments representing the training configuration.
+                method uses a default trainer. Defaults to None.
+            **kwargs (dict): Arbitrary keyword arguments representing the training configuration. These arguments are
+                used to customize various aspects of the training process.
+
+        Returns:
+            (dict | None): Training metrics if available and training is successful; otherwise, None.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
+            PermissionError: If there is a permission issue with the HUB session.
+            ModuleNotFoundError: If the HUB SDK is not installed.
         """
         self._check_is_pytorch_model()
         if hasattr(self.session, "model") and self.session.model.id:  # Ultralytics HUB session with loaded model
@@ -399,10 +609,24 @@ class Model(nn.Module):
 
     def tune(self, use_ray=False, iterations=10, *args, **kwargs):
         """
-        Runs hyperparameter tuning, optionally using Ray Tune. See ultralytics.utils.tuner.run_ray_tune for Args.
+        Conducts hyperparameter tuning for the model, with an option to use Ray Tune.
+
+        This method supports two modes of hyperparameter tuning: using Ray Tune or a custom tuning method.
+        When Ray Tune is enabled, it leverages the 'run_ray_tune' function from the ultralytics.utils.tuner module.
+        Otherwise, it uses the internal 'Tuner' class for tuning. The method combines default, overridden, and
+        custom arguments to configure the tuning process.
+
+        Args:
+            use_ray (bool): If True, uses Ray Tune for hyperparameter tuning. Defaults to False.
+            iterations (int): The number of tuning iterations to perform. Defaults to 10.
+            *args (list): Variable length argument list for additional arguments.
+            **kwargs (dict): Arbitrary keyword arguments. These are combined with the model's overrides and defaults.
 
         Returns:
             (dict): A dictionary containing the results of the hyperparameter search.
+
+        Raises:
+            AssertionError: If the model is not a PyTorch model.
         """
         self._check_is_pytorch_model()
         if use_ray:
@@ -426,31 +650,81 @@ class Model(nn.Module):
 
     @property
     def names(self):
-        """Returns class names of the loaded model."""
+        """
+        Retrieves the class names associated with the loaded model.
+
+        This property returns the class names if they are defined in the model. It checks the class names for validity
+        using the 'check_class_names' function from the ultralytics.nn.autobackend module.
+
+        Returns:
+            (list | None): The class names of the model if available, otherwise None.
+        """
         from ultralytics.nn.autobackend import check_class_names
 
         return check_class_names(self.model.names) if hasattr(self.model, "names") else None
 
     @property
     def device(self):
-        """Returns device if PyTorch model."""
+        """
+        Retrieves the device on which the model's parameters are allocated.
+
+        This property is used to determine whether the model's parameters are on CPU or GPU. It only applies to models
+        that are instances of nn.Module.
+
+        Returns:
+            (torch.device | None): The device (CPU/GPU) of the model if it is a PyTorch model, otherwise None.
+        """
         return next(self.model.parameters()).device if isinstance(self.model, nn.Module) else None
 
     @property
     def transforms(self):
-        """Returns transform of the loaded model."""
+        """
+        Retrieves the transformations applied to the input data of the loaded model.
+
+        This property returns the transformations if they are defined in the model.
+
+        Returns:
+            (object | None): The transform object of the model if available, otherwise None.
+        """
         return self.model.transforms if hasattr(self.model, "transforms") else None
 
     def add_callback(self, event: str, func):
-        """Add a callback."""
+        """
+        Adds a callback function for a specified event.
+
+        This method allows the user to register a custom callback function that is triggered on a specific event during
+        model training or inference.
+
+        Args:
+            event (str): The name of the event to attach the callback to.
+            func (callable): The callback function to be registered.
+
+        Raises:
+            ValueError: If the event name is not recognized.
+        """
         self.callbacks[event].append(func)
 
     def clear_callback(self, event: str):
-        """Clear all event callbacks."""
+        """
+        Clears all callback functions registered for a specified event.
+
+        This method removes all custom and default callback functions associated with the given event.
+
+        Args:
+            event (str): The name of the event for which to clear the callbacks.
+
+        Raises:
+            ValueError: If the event name is not recognized.
+        """
         self.callbacks[event] = []
 
     def reset_callbacks(self):
-        """Reset all registered callbacks."""
+        """
+        Resets all callbacks to their default functions.
+
+        This method reinstates the default callback functions for all events, removing any custom callbacks that were
+        added previously.
+        """
         for event in callbacks.default_callbacks.keys():
             self.callbacks[event] = [callbacks.default_callbacks[event][0]]
 

ultralytics/nn/tasks.py

@@ -631,7 +631,7 @@ def torch_safe_load(weight):
                 "ultralytics.yolo.data": "ultralytics.data",
             }
         ):  # for legacy 8.0 Classify and Pose models
-            return torch.load(file, map_location="cpu"), file  # load
+            ckpt = torch.load(file, map_location="cpu")
 
     except ModuleNotFoundError as e:  # e.name is missing module name
         if e.name == "models":
@@ -651,8 +651,17 @@ def torch_safe_load(weight):
             f"run a command with an official YOLOv8 model, i.e. 'yolo predict model=yolov8n.pt'"
         )
         check_requirements(e.name)  # install missing module
+        ckpt = torch.load(file, map_location="cpu")
 
-        return torch.load(file, map_location="cpu"), file  # load
+    if not isinstance(ckpt, dict):
+        # File is likely a YOLO instance saved with i.e. torch.save(model, "saved_model.pt")
+        LOGGER.warning(
+            f"WARNING ⚠️ The file '{weight}' appears to be improperly saved or formatted. "
+            f"For optimal results, use model.save('filename.pt') to correctly save YOLO models."
+        )
+        ckpt = {"model": ckpt.model}
+
+    return ckpt, file  # load
 
 
 def attempt_load_weights(weights, device=None, inplace=True, fuse=False):