Merge pull request #772 from ACEsuit/select_head_fix

Fix default outpout_file in select_head, and add argument to list heads

mace/cli/select_head.py

@@ -7,26 +7,47 @@
 
 def main():
     parser = ArgumentParser()
-    parser.add_argument(
+    grp = parser.add_mutually_exclusive_group()
+    grp.add_argument(
         "--head_name",
         "-n",
         help="name of the head to extract",
         default=None,
     )
+    grp.add_argument(
+        "--list_heads",
+        "-l",
+        action="store_true",
+        help="list names of the heads",
+    )
+    parser.add_argument(
+        "--target_device",
+        "-d",
+        help="target device, defaults to model's current device",
+    )
     parser.add_argument(
         "--output_file",
         "-o",
-        help="name for output model, defaults to model_file.target_device",
+        help="name for output model, defaults to model.head_name, followed by .target_device if specified",
     )
     parser.add_argument("model_file", help="input model file path")
     args = parser.parse_args()
 
-    if args.output_file is None:
-        args.output_file = args.model_file + "." + args.target_device
-
     model = torch.load(args.model_file)
-    model_single = remove_pt_head(model, args.head_name)
-    torch.save(model_single, args.output_file)
+
+    if args.list_heads:
+        print("Available heads:")
+        print("\n".join(["  " + h for h in model.heads]))
+    else:
+
+        if args.output_file is None:
+            args.output_file = args.model_file + "." + args.head_name + ("." + args.target_device if (args.target_device is not None) else "")
+
+        model_single = remove_pt_head(model, args.head_name)
+        if args.target_device is not None:
+            target_device = str(next(model.parameters()).device)
+            model_single.to(target_device)
+        torch.save(model_single, args.output_file)
 
 
 if __name__ == "__main__":