diff --git a/mbtrack2/utilities/beamloading.py b/mbtrack2/utilities/beamloading.py
index 3dbc7d53f85afad83550d408e45d02af08163cbb..84a23a18645ba3e0bcbd8e50a2797ae2c3f3c115 100644
--- a/mbtrack2/utilities/beamloading.py
+++ b/mbtrack2/utilities/beamloading.py
@@ -138,7 +138,7 @@ class BeamLoadingEquilibrium():
                 cavity = self.cavity_list[i]
                 if cavity.m == 1:
                     cavity.theta = x[-1]
-                    cavity.set_generator(0.5)
+                    cavity.set_generator(self.I0)
                     self.update_potentials()
         else:
             self.F = x[::2]
@@ -180,8 +180,8 @@ class BeamLoadingEquilibrium():
         if z2 is None:
             z2 = self.B2
         z0 = np.linspace(z1, z2, 1000)
-        plt.plot(z0, self.rho(z0))
-        plt.xlabel("z [m]")
+        plt.plot(z0/c*1e12, self.rho(z0))
+        plt.xlabel(r"$\tau$ [ps]")
         plt.title("Equilibrium bunch profile")
         
     def voltage(self, z):