from #928

nyu-mll · Oct 8, 2019 · 7f4c3c0 · 7f4c3c0
1 parent c87598f
commit 7f4c3c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jiant/trainer.py b/jiant/trainer.py
@@ -1204,7 +1204,7 @@ def _restore_checkpoint(self, phase, tasks=None):
             self._serialization_dir, task_directory, "_".join(["metric", suffix])
         )
 
-        model_state = torch.load(model_path, map_location=device_mapping(self._cuda_device))
+        model_state = torch.load(model_path)
 
         for name, param in self._model.named_parameters():
             if param.requires_grad and name not in model_state: