From b7542b63c4643aec3cff43cbec5362a385ecb921 Mon Sep 17 00:00:00 2001
From: Francesco Innocenti <F.Innocenti@sussex.ac.uk>
Date: Wed, 27 Nov 2024 16:09:47 +0000
Subject: [PATCH] Rename updates' outputs for clarity.

---
 jpc/_core/_updates.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/jpc/_core/_updates.py b/jpc/_core/_updates.py
index 1f6f51e..f3a1aea 100644
--- 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
     }