add cpu support via map_location in test.py

Author: Kirili4ik

test.py

@@ -18,7 +18,10 @@
 
 def main(config, out_file):
     logger = config.get_logger("test")
-
+    
+    # define cpu or gpu if possible
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
     # text_encoder
     text_encoder = CTCCharTextEncoder.get_simple_alphabet()
 
@@ -30,14 +33,13 @@ def main(config, out_file):
     logger.info(model)
 
     logger.info("Loading checkpoint: {} ...".format(config.resume))
-    checkpoint = torch.load(config.resume)
+    checkpoint = torch.load(config.resume, map_location=device)
     state_dict = checkpoint["state_dict"]
     if config["n_gpu"] > 1:
         model = torch.nn.DataParallel(model)
     model.load_state_dict(state_dict)
 
     # prepare model for testing
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = model.to(device)
     model.eval()