Fix init

Lightning-AI · May 23, 2024 · d882a8f · d882a8f
1 parent 2174054
commit d882a8f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/litgpt/utils.py b/litgpt/utils.py
@@ -514,5 +514,7 @@ def instantiate_torch_optimizer(optimizer, model_parameters, **kwargs):
         optimizer_cls = getattr(torch.optim, optimizer)
         optimizer = optimizer_cls(model_parameters, **kwargs)
     else:
-        optimizer = instantiate_class(model_parameters, optimizer, **kwargs)
+        optimizer = dict(optimizer)  # copy
+        optimizer["init_args"].update(kwargs)
+        optimizer = instantiate_class(model_parameters, optimizer)
     return optimizer