Skip to content

Commit

Permalink
Add map_location to load_state function
Browse files Browse the repository at this point in the history
  • Loading branch information
[email protected] authored and sidharthms committed Apr 20, 2020
1 parent 5fc2cd4 commit 9123fa3
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions deepmatcher/models/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -458,13 +458,13 @@ def save_state(self, path, include_meta=True):
state[k] = getattr(self, k)
torch.save(state, path, pickle_module=dill)

def load_state(self, path):
def load_state(self, path, map_location=None):
r"""Load the model state from a file in a certain path.
Args:
path (string): The path to load the model state from.
"""
state = torch.load(path, pickle_module=dill)
state = torch.load(path, pickle_module=dill, map_location=map_location)
for k, v in six.iteritems(state):
if k != 'model':
self._train_buffers.add(k)
Expand Down

0 comments on commit 9123fa3

Please sign in to comment.