corrected gamma in the case not updated

google · Oct 7, 2022 · 90dbb5b · 90dbb5b
1 parent 2ab6b2d
commit 90dbb5b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jaxopt/_src/lbfgs.py b/jaxopt/_src/lbfgs.py
@@ -142,7 +142,7 @@ class LbfgsState(NamedTuple):
   s_history: Any
   y_history: Any
   rho_history: jnp.ndarray
-  gamma: float = 1.0
+  gamma: Any = jnp.array(1.0)
   aux: Optional[Any] = None
 
 
@@ -349,7 +349,7 @@ def update(self,
     if self.use_gamma:
       gamma = compute_gamma(s_history, y_history, last)
     else:
-      gamma = 1.0
+      gamma = jnp.array(1.0)
 
     new_state = LbfgsState(iter_num=state.iter_num + 1,
                            value=new_value,