Merge pull request facebookresearch#144 from facebookresearch/parallel

fix for saving models trained with DataParallel
tunde-ajayi · May 17, 2018 · 50d0e49 · 50d0e49
2 parents 3b19410 + c3a19f8
commit 50d0e49
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/drqa/reader/model.py b/drqa/reader/model.py
@@ -394,7 +394,11 @@ def decode_candidates(score_s, score_e, candidates, top_n=1, max_len=None):
     # --------------------------------------------------------------------------
 
     def save(self, filename):
-        state_dict = copy.copy(self.network.state_dict())
+        if self.parallel:
+            network = self.network.module
+        else:
+            network = self.network
+        state_dict = copy.copy(network.state_dict())
         if 'fixed_embedding' in state_dict:
             state_dict.pop('fixed_embedding')
         params = {
@@ -409,8 +413,12 @@ def save(self, filename):
             logger.warning('WARN: Saving failed... continuing anyway.')
 
     def checkpoint(self, filename, epoch):
+        if self.parallel:
+            network = self.network.module
+        else:
+            network = self.network
         params = {
-            'state_dict': self.network.state_dict(),
+            'state_dict': network.state_dict(),
             'word_dict': self.word_dict,
             'feature_dict': self.feature_dict,
             'args': self.args,