torch.save

torch.save(obj, f, pickle_module=<module 'pickle' from '/opt/conda/lib/python3.6/pickle.py'>, pickle_protocol=2, _use_new_zipfile_serialization=True)

Saves an object to a disk file.

See also: Recommended approach for saving a model

Parameters

• obj – saved object

• f – a file-like object (has to implement write and flush) or a string containing a file name

• pickle_module – module used for pickling metadata and objects

• pickle_protocol – can be specified to override the default protocol

Note

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

Warning

If you are using Python 2, torch.save() does NOT support StringIO.StringIO as a valid file-like object. This is because the write method should return the number of bytes written; StringIO.write() does not do this.

Please use something like io.BytesIO instead.

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)