fix: Fix gradients plot

plotscripts.py

@@ -145,9 +145,9 @@ def plot_gradients(
     """
     grads = dict(np.load(path / f"{GRADS_FILE + '.npz'}"))
     config = json.loads((path / OPTIMIZATION_FILE).read_text())
-
     ave_grads = []
-
+    print(len(grads["0"]))
+    # print([i for i in ])
     for epoch in grads:
         for grads_list in grads[epoch]:
             ave_grads.append(np.mean(np.abs(grads_list)))
@@ -162,7 +162,7 @@ def plot_gradients(
         label=r"$\langle |\partial_{\theta_i}\text{L}| \rangle_i$",
     )
     for b in range(config["nboost"] - 1):
-        boost_x = config["boost_frequency"] * (b + 1)
+        boost_x = len(grads[str(b)]) * (b + 1)
         if b == 0:
             plt.plot(
                 (boost_x, boost_x + 1),