torch.save (original) (raw)

torch.save(obj, f, pickle_module=pickle, pickle_protocol=2, _use_new_zipfile_serialization=True)[source]#

Saves an object to a disk file.

See also: Saving and loading tensors

See Layout Control for more advanced tools to manipulate a checkpoint.

Parameters

• obj (object) – saved object
• f (Union[_str,_ PathLike_[_str] , IO_[_bytes] ]) – a file-like object (has to implement write and flush) or a string or os.PathLike object containing a file name
• pickle_module (Any) – module used for pickling metadata and objects
• pickle_protocol (int) – can be specified to override the default protocol

Note

A common PyTorch convention is to save tensors using .pt file extension.

Note

The 1.6 release of PyTorch switched torch.save to use a new zipfile-based file format. torch.load still retains the ability to load files in the old format. If for any reason you want torch.saveto use the old format, pass the kwarg _use_new_zipfile_serialization=False.

Example

Save to file

x = torch.tensor([0, 1, 2, 3, 4]) torch.save(x, "tensor.pt")

Save to io.BytesIO buffer

buffer = io.BytesIO() torch.save(x, buffer)