Apply suggestion from @Copilot

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

Author: xiu-cs

animals/demo/vis_animals.py

@@ -339,7 +339,7 @@ def get_pose3D(path, output_dir, type='image'):
     model_dict = model['CFM'].state_dict()
     model_path = args.saved_model_path
     print(f"Loading model from: {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)