Fix bug in factory

Signed-off-by: ashors1 <[email protected]>
NVIDIA · Jul 18, 2024 · 71e36fb · 71e36fb
1 parent 0147133
commit 71e36fb
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/nemo/collections/llm/utils.py b/nemo/collections/llm/utils.py
@@ -42,9 +42,8 @@ def factory(*args: Any, **kwargs: Any) -> Union[Callable[[T], T], T]:
     try:
         import nemo_sdk as sdk
 
-        if not args and not kwargs:
-            # Used as @factory without arguments
-            return sdk.factory()
+        if not args:
+            return sdk.factory(**kwargs)
         else:
             # Used as @factory(*args, **kwargs)
             return sdk.factory(*args, **kwargs)