torch.jit.load¶
-
torch.jit.
load
(f, map_location=None, _extra_files=None)[source]¶ Load a
ScriptModule
orScriptFunction
previously saved withtorch.jit.save
All previously saved modules, no matter their device, are first loaded onto CPU, and then are moved to the devices they were saved from. If this fails (e.g. because the run time system doesn’t have certain devices), an exception is raised.
- Parameters
f – a file-like object (has to implement read, readline, tell, and seek), or a string containing a file name
map_location (string or torch.device) – A simplified version of
map_location
in torch.jit.save used to dynamically remap storages to an alternative set of devices._extra_files (dictionary of filename to content) – The extra filenames given in the map would be loaded and their content would be stored in the provided map.
- Returns
A
ScriptModule
object.
Example:
import torch import io torch.jit.load('scriptmodule.pt') # Load ScriptModule from io.BytesIO object with open('scriptmodule.pt', 'rb') as f: buffer = io.BytesIO(f.read()) # Load all tensors to the original device torch.jit.load(buffer) # Load all tensors onto CPU, using a device buffer.seek(0) torch.jit.load(buffer, map_location=torch.device('cpu')) # Load all tensors onto CPU, using a string buffer.seek(0) torch.jit.load(buffer, map_location='cpu') # Load with extra files. extra_files = {'foo.txt': ''} # values will be replaced with data torch.jit.load('scriptmodule.pt', _extra_files=extra_files) print(extra_files['foo.txt'])