From 6de35365cf94b4fa38939425f035961becd2177a Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Mon, 11 Nov 2024 15:55:50 -0700
Subject: [PATCH] Fix/add some type annotations

---
 torchtune/config/_instantiate.py | 8 ++++----
 torchtune/config/_parse.py       | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/torchtune/config/_instantiate.py b/torchtune/config/_instantiate.py
index 279fe22629..8ebc1b369d 100644
--- a/torchtune/config/_instantiate.py
+++ b/torchtune/config/_instantiate.py
@@ -18,11 +18,11 @@ def _create_component(
     _component_: Callable[..., Any],
     args: Tuple[Any, ...],
     kwargs: Dict[str, Any],
-):
+) -> Any:
     return _component_(*args, **kwargs)
 
 
-def _instantiate_node(node: Dict[str, Any], *args: Tuple[Any, ...]):
+def _instantiate_node(node: Dict[str, Any], *args: Any) -> Any:
     """
     Creates the object specified in _component_ field with provided positional args
     and kwargs already merged. Raises an InstantiationError if _component_ is not specified.
@@ -40,8 +40,8 @@ def _instantiate_node(node: Dict[str, Any], *args: Tuple[Any, ...]):
 
 def instantiate(
     config: DictConfig,
-    *args: Tuple[Any, ...],
-    **kwargs: Dict[str, Any],
+    *args: Any,
+    **kwargs: Any,
 ) -> Any:
     """
     Given a DictConfig with a _component_ field specifying the object to instantiate and
diff --git a/torchtune/config/_parse.py b/torchtune/config/_parse.py
index 8472094f0b..f7ddf63489 100644
--- a/torchtune/config/_parse.py
+++ b/torchtune/config/_parse.py
@@ -65,7 +65,7 @@ def parse_known_args(self, *args, **kwargs) -> Tuple[Namespace, List[str]]:
         return namespace, unknown_args
 
 
-def parse(recipe_main: Recipe) -> Callable[[Recipe], Any]:
+def parse(recipe_main: Recipe) -> Callable[[], Any]:
     """
     Decorator that handles parsing the config file and CLI overrides
     for a recipe. Use it on the recipe's main function.