Remove dill package from Ultralytics (#16574)
Signed-off-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
parent
060214d066
commit
c327b0aae1
5 changed files with 5 additions and 18 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue