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

mace/cli/select_head.py

@@ -7,12 +7,19 @@
 
 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(
         "--output_file",
         "-o",
@@ -21,12 +28,15 @@ def main():
     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 + "." + str(next(model.parameters()).device)
+        model_single = remove_pt_head(model, args.head_name)
+        torch.save(model_single, args.output_file)
 
 
 if __name__ == "__main__":