Add patch to assistant

Lightning-AI · Dec 9, 2024 · ce4763d · ce4763d
1 parent c600bd3
commit ce4763d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/.actions/assistant.py b/.actions/assistant.py
@@ -483,14 +483,15 @@ def convert_version2nightly(ver_file: str = "src/version.info") -> None:
 
 
 if __name__ == "__main__":
-    import jsonargparse
     import sys
 
-    def _parse_known_args_patch(self, args: Any=None, namespace: Any=None) -> tuple[Any]:
+    import jsonargparse
+
+    def _parse_known_args_patch(self, args: Any = None, namespace: Any = None) -> tuple[Any]:
         namespace, args = super(jsonargparse.ArgumentParser, self)._parse_known_args(args, namespace, intermixed=False)
         return namespace, args
 
     if sys.version_info >= (3, 12, 8):
-        setattr(ArgumentParser, "_parse_known_args", _parse_known_args_patch)
+        setattr(jsonargparse.ArgumentParser, "_parse_known_args", _parse_known_args_patch)
 
     jsonargparse.CLI(AssistantCLI, as_positional=False)
diff --git a/src/lightning/pytorch/cli.py b/src/lightning/pytorch/cli.py
@@ -48,7 +48,7 @@
         set_config_read_mode,
     )
 
-    def _parse_known_args_patch(self, args: Any=None, namespace: Any=None) -> tuple[Any]:
+    def _parse_known_args_patch(self, args: Any = None, namespace: Any = None) -> tuple[Any]:
         namespace, args = super(ArgumentParser, self)._parse_known_args(args, namespace, intermixed=False)
         return namespace, args