Rename updates' outputs for clarity.

thebuckleylab · Nov 27, 2024 · b7542b6 · b7542b6
1 parent 6115c57
commit b7542b6
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/jpc/_core/_updates.py b/jpc/_core/_updates.py
@@ -52,8 +52,8 @@ def update_activities(
         "energy": energy,
         "activities": activities,
         "activity_grads": activity_grads,
-        "activity_optim": optim,
-        "activity_opt_state": opt_state
+        "optim": optim,
+        "opt_state": opt_state
     }
 
 
@@ -102,6 +102,6 @@ def update_params(
         "model": model,
         "skip_model": skip_model,
         "param_grads": param_grads,
-        "param_optim": optim,
-        "param_opt_state": opt_state
+        "optim": optim,
+        "opt_state": opt_state
     }