Skip to content

Commit

Permalink
Leave native save as False (#2138)
Browse files Browse the repository at this point in the history
* Custom objects are not saved using saftensors

* Leave save as false
  • Loading branch information
muellerzr committed Nov 9, 2023
1 parent 9674b40 commit 6727ac4
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions src/accelerate/utils/other.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ def wait_for_everyone():
PartialState().wait_for_everyone()


def save(obj, f, save_on_each_node: bool = False, safe_serialization: bool = True):
def save(obj, f, save_on_each_node: bool = False, safe_serialization: bool = False):
"""
Save the data to disk. Use in place of `torch.save()`.
Expand All @@ -126,8 +126,8 @@ def save(obj, f, save_on_each_node: bool = False, safe_serialization: bool = Tru
The file (or file-like object) to use to save the data
save_on_each_node (`bool`, *optional*, defaults to `False`):
Whether to only save on the global main process
safe_serialization (`bool`, *optional*, defaults to `True`):
Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
safe_serialization (`bool`, *optional*, defaults to `False`):
Whether to save `obj` using `safetensors` or the traditional PyTorch way (that uses `pickle`).
"""
save_func = torch.save if not safe_serialization else partial(safe_save_file, metadata={"format": "pt"})
if PartialState().distributed_type == DistributedType.TPU:
Expand Down

0 comments on commit 6727ac4

Please sign in to comment.