Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: xiu-cs

demo/vis_in_the_wild.py

@@ -287,7 +287,7 @@ def get_pose3D(path, output_dir, type='image'):
     model_dict = model['CFM'].state_dict()
     model_path = args.model_weights_path
     print(model_path)
-    pre_dict = torch.load(model_path, weights_only=True)
+    pre_dict = torch.load(model_path, map_location=device, weights_only=True)
     for name, key in model_dict.items():
         model_dict[name] = pre_dict[name]
     model['CFM'].load_state_dict(model_dict)