diff --git a/graph_net/torch/test_target_device.py b/graph_net/torch/test_target_device.py
index ee46ceee6..5a4b0daf5 100644
--- a/graph_net/torch/test_target_device.py
+++ b/graph_net/torch/test_target_device.py
@@ -65,12 +65,12 @@ def test_single_model(args):
     eval_backend_perf.eval_single_model_with_single_backend(eval_args)
 
     ref_dump = utils.get_output_path(args.reference_dir, args.model_path)
-    ref_out = torch.load(str(ref_dump))
+    ref_out = torch.load(str(ref_dump), map_location=torch.device(args.device))
     ref_log = utils.get_log_path(args.reference_dir, args.model_path)
     ref_time_stats = eval_backend_diff.parse_time_stats_from_reference_log(ref_log)
 
     target_dump = utils.get_output_path(target_dir, args.model_path)
-    target_out = torch.load(str(target_dump))
+    target_out = torch.load(str(target_dump), map_location=torch.device(args.device))
     target_log = utils.get_log_path(target_dir, args.model_path)
     target_time_stats = eval_backend_diff.parse_time_stats_from_reference_log(
         target_log
@@ -99,7 +99,7 @@ def test_multi_models(args):
 
         if test_samples is None or os.path.abspath(model_path) in test_samples:
             print(
-                f"[{sample_idx}] {module_name}, model_path: {model_path}",
+                f"[{sample_idx}][Processing] {module_name}, model_path: {model_path}",
                 file=sys.stderr,
                 flush=True,
             )