Update optimization.py

YaoGroup · Aug 6, 2024 · 78b6931 · 78b6931
1 parent ca6e840
commit 78b6931
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/diffice_jax/optimizer/optimization.py b/diffice_jax/optimizer/optimization.py
@@ -1,3 +1,7 @@
+"""
+@author: Yongji Wang
+"""
+
 import sys
 import jax.numpy as jnp
 import optax
@@ -138,4 +142,4 @@ def lbfgs_optimizer(lossf, params, data, epoch):
     num_iter = results.num_objective_evaluations
     loss_all = func_lbfgs.loss
     print(f" Total iterations: {num_iter}")
-    return params, loss_all
+    return params, loss_all