updated torch.load params

Author: bsh98

model_coverage.py

@@ -91,7 +91,7 @@ def mask_text(self, text_tokenized):
         return masked
 
     def reload_model(self, model_file):
-        print(self.model.load_state_dict(torch.load(model_file), strict=False))
+        print(self.model.load_state_dict(torch.load(model_file, map_location=torch.device(self.device)), strict=False))
 
     def save_model(self, model_file):
         torch.save(self.model.state_dict(), model_file)

model_generator.py

@@ -36,7 +36,7 @@ def __init__(self, max_output_length=25, max_input_length=300, device='cpu', tok
         self.mode = "train"
 
     def reload(self, from_file):
-        print(self.model.load_state_dict(torch.load(from_file), strict=False))
+        print(self.model.load_state_dict(torch.load(from_file, map_location=torch.device(self.device)), strict=False))
 
     def save(self, to_file):
         torch.save(self.model.state_dict(), to_file)