More elegant solution

metatensor · Nov 14, 2024 · 8cf49c4 · 8cf49c4
1 parent ce06288
commit 8cf49c4
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/metatrain/cli/export.py b/src/metatrain/cli/export.py
@@ -72,8 +72,11 @@ def _prepare_export_model_args(args: argparse.Namespace) -> None:
         architecture_name=architecture_name,
         **args.__dict__,
     )
-    if "huggingface_api_token" in args.__dict__:
-        args.__dict__.pop("huggingface_api_token")
+    keys_to_keep = ["model", "output"]  # only these are needed for `export_model``
+    original_keys = list(args.__dict__.keys())
+    for key in original_keys:
+        if key not in keys_to_keep:
+            args.__dict__.pop(key)
 
 
 def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> None: