diff --git a/assets/Notebooks/VR_exercise.ipynb b/assets/Notebooks/VR_exercise.ipynb
index c8207c61..45167ebc 100644
--- a/assets/Notebooks/VR_exercise.ipynb
+++ b/assets/Notebooks/VR_exercise.ipynb
@@ -339,6 +339,7 @@
         "    flops_trajectory = [flops]\n",
         "    for _ in range(iterations // epoch_length):\n",
         "        full_grad = grad(lambda w: least_squares_loss(w, X, y, mu))(x)\n",
+        "        flops += 2 * m * n
         "        x_tilde = x\n",
         "        for _ in range(epoch_length):\n",
         "            indices = np.random.choice(m, batch_size, replace=False)\n",