torch.jit.load#
- torch.jit.load(f,map_location=None,_extra_files=None,_restore_shapes=False)[source]#
Load a
ScriptModuleorScriptFunctionpreviously 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 israised.
- Parameters
f – a file-like object (has to implement read, readline, tell, and seek),or a string containing a file name
map_location (string ortorch.device) – A simplified version of
map_locationintorch.jit.save used to dynamically remapstorages to an alternative set of devices._extra_files (dictionary offilename to content) – The extrafilenames given in the map would be loaded and their contentwould be stored in the provided map.
_restore_shapes (bool) – Whether or not to retrace the module on load using stored inputs
- Returns
A
ScriptModuleobject.
Warning
It is possible to construct malicious pickle data which will execute arbitrary codeduring func:torch.jit.load. Never load data that could have come from an untrustedsource, or that could have been tampered with.Only load data you trust.
Example:.. testcode:
importtorchimportiotorch.jit.load('scriptmodule.pt')# Load ScriptModule from io.BytesIO objectwithopen('scriptmodule.pt','rb')asf:buffer=io.BytesIO(f.read())# Load all tensors to the original devicetorch.jit.load(buffer)# Load all tensors onto CPU, using a devicebuffer.seek(0)torch.jit.load(buffer,map_location=torch.device('cpu'))# Load all tensors onto CPU, using a stringbuffer.seek(0)torch.jit.load(buffer,map_location='cpu')# Load with extra files.extra_files={'foo.txt':''}# values will be replaced with datatorch.jit.load('scriptmodule.pt',_extra_files=extra_files)print(extra_files['foo.txt'])