fix Making a prediction on CPU after training on GPU google#89

uisrnn/uisrnn.py

@@ -151,7 +151,7 @@ def load(self, filepath):
     Args:
       filepath: the path of the file.
     """
-    var_dict = torch.load(filepath)
+    var_dict = torch.load(filepath, map_location=self.device)
     self.rnn_model.load_state_dict(var_dict['rnn_state_dict'])
     self.rnn_init_hidden = nn.Parameter(
         torch.from_numpy(var_dict['rnn_init_hidden']).to(self.device))