From c327b0aae17f84c6f587c2ba5653366bc0165945 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 1 Oct 2024 10:10:26 +0200
Subject: [PATCH] Remove `dill` package from Ultralytics (#16574)

Signed-off-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
---
 ultralytics/engine/model.py      |  5 ++---
 ultralytics/models/sam/build.py  |  2 --
 ultralytics/utils/files.py       |  2 +-
 ultralytics/utils/patches.py     | 12 +-----------
 ultralytics/utils/torch_utils.py |  2 +-
 5 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py
index fcafc9f78d..c4db53426a 100644
--- a/ultralytics/engine/model.py
+++ b/ultralytics/engine/model.py
@@ -377,7 +377,7 @@ class Model(nn.Module):
         self.model.load(weights)
         return self
 
-    def save(self, filename: Union[str, Path] = "saved_model.pt", use_dill=True) -> None:
+    def save(self, filename: Union[str, Path] = "saved_model.pt") -> None:
         """
         Saves the current model state to a file.
 
@@ -386,7 +386,6 @@ class Model(nn.Module):
 
         Args:
             filename (Union[str, Path]): The name of the file to save the model to.
-            use_dill (bool): Whether to try using dill for serialization if available.
 
         Raises:
             AssertionError: If the model is not a PyTorch model.
@@ -408,7 +407,7 @@ class Model(nn.Module):
             "license": "AGPL-3.0 License (https://ultralytics.com/license)",
             "docs": "https://docs.ultralytics.com",
         }
-        torch.save({**self.ckpt, **updates}, filename, use_dill=use_dill)
+        torch.save({**self.ckpt, **updates}, filename)
 
     def info(self, detailed: bool = False, verbose: bool = True):
         """
diff --git a/ultralytics/models/sam/build.py b/ultralytics/models/sam/build.py
index 0e7ddedcf0..e110531244 100644
--- a/ultralytics/models/sam/build.py
+++ b/ultralytics/models/sam/build.py
@@ -210,8 +210,6 @@ def _build_sam(
             state_dict = torch.load(f)
         sam.load_state_dict(state_dict)
     sam.eval()
-    # sam.load_state_dict(torch.load(checkpoint), strict=True)
-    # sam.eval()
     return sam
 
 
diff --git a/ultralytics/utils/files.py b/ultralytics/utils/files.py
index d0953c748e..29c68d48de 100644
--- a/ultralytics/utils/files.py
+++ b/ultralytics/utils/files.py
@@ -219,4 +219,4 @@ def update_models(model_names=("yolov8n.pt",), source_dir=Path("."), update_name
 
         # Save model using model.save()
         print(f"Re-saving {model_name} model to {save_path}")
-        model.save(save_path, use_dill=False)
+        model.save(save_path)
diff --git a/ultralytics/utils/patches.py b/ultralytics/utils/patches.py
index d918e0efea..e9ba5dfb30 100644
--- a/ultralytics/utils/patches.py
+++ b/ultralytics/utils/patches.py
@@ -86,25 +86,15 @@ def torch_load(*args, **kwargs):
     return _torch_load(*args, **kwargs)
 
 
-def torch_save(*args, use_dill=True, **kwargs):
+def torch_save(*args, **kwargs):
     """
     Optionally use dill to serialize lambda functions where pickle does not, adding robustness with 3 retries and
     exponential standoff in case of save failure.
 
     Args:
         *args (tuple): Positional arguments to pass to torch.save.
-        use_dill (bool): Whether to try using dill for serialization if available. Defaults to True.
         **kwargs (Any): Keyword arguments to pass to torch.save.
     """
-    try:
-        assert use_dill
-        import dill as pickle
-    except (AssertionError, ImportError):
-        import pickle
-
-    if "pickle_module" not in kwargs:
-        kwargs["pickle_module"] = pickle
-
     for i in range(4):  # 3 retries
         try:
             return _torch_save(*args, **kwargs)
diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py
index e7fcca0ad7..5c477fb1e9 100644
--- a/ultralytics/utils/torch_utils.py
+++ b/ultralytics/utils/torch_utils.py
@@ -595,7 +595,7 @@ def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "", updates: dict
 
     # Save
     combined = {**metadata, **x, **(updates or {})}
-    torch.save(combined, s or f, use_dill=False)  # combine dicts (prefer to the right)
+    torch.save(combined, s or f)  # combine dicts (prefer to the right)
     mb = os.path.getsize(s or f) / 1e6  # file size
     LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
     return combined