Accept arguments in the optimizer advances #17

pab1s · Jun 24, 2024 · 125cd7e · 125cd7e
1 parent d4c8ba2
commit 125cd7e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/factories/loss_factory.py b/factories/loss_factory.py
@@ -4,5 +4,5 @@
 class LossFactory(Factory):
     def __init__(self):
         super().__init__()
-        self.register("CrossEntropyLoss", lambda: CrossEntropyLoss())
-        self.register("MSELoss", lambda: MSELoss())
+        self.register("CrossEntropyLoss", lambda **kwargs: CrossEntropyLoss(**kwargs))
+        self.register("MSELoss", lambda **kwargs: MSELoss(**kwargs))