diff --git a/dmff/admp/pme.py b/dmff/admp/pme.py
index 21ea5d3a1..14b9d4056 100755
--- a/dmff/admp/pme.py
+++ b/dmff/admp/pme.py
@@ -209,6 +209,7 @@ def get_energy(
                     U_init=U_init * 10.0,
                     steps_pol=self.steps_pol,
                 )  # nm to angstrom
+                self.U_ind = U_ind
                 # here we rely on Feynman-Hellman theorem, drop the term dV/dU*dU/dr !
                 # self.U_ind = jax.lax.stop_gradient(U_ind)
                 energy = energy_fn(