Added flag to allow loading of CUDA-trained parameters to CPU.
@@ -181,12 +181,12 @@ class ParamStoreDict(object):
Save parameters to disk
:param filename: file name to save to
- :type name: str
+ :type filename: str
"""
with open(filename, "wb") as output_file:
torch.save(self.get_state(), output_file)
- def load(self, filename):
+ def load(self, filename, map_location=None):
"""
Loads parameters from disk
@@ -199,10 +199,12 @@ class ParamStoreDict(object):
pyro.module('module', nn, update_module_params=True)
:param filename: file name to load from
- :type name: str
+ :type filename: str
+ :param map_location: specifies how to remap storage locations
+ :type map_location: function, torch.device, string or a dict
"""
with open(filename, "rb") as input_file:
- state = torch.load(input_file)
+ state = torch.load(input_file, map_location)
self.set_state(state)
