diff --git a/nuclear-qc/vqe/qng.ipynb b/nuclear-qc/vqe/qng.ipynb
index 7f69b73..aa78516 100644
--- a/nuclear-qc/vqe/qng.ipynb
+++ b/nuclear-qc/vqe/qng.ipynb
@@ -19,7 +19,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -63,36 +63,52 @@
     "dev = qml.device(\"default.qubit\", wires=3)\n",
     "\n",
     "# Defining our ansatz circuit for the N=3 case\n",
-    "\n",
-    "def circuit(params,wires):\n",
+    "def circuit(params, wires):\n",
     "    t0 = params[0]\n",
     "    t1 = params[1]\n",
     "    qml.PauliX(wires=0)\n",
     "    qml.RY(t1, wires=1)\n",
     "    qml.RY(t0, wires=2)\n",
-    "    qml.CNOT(wires=[2,0])\n",
-    "    qml.CNOT(wires=[0,1])\n",
+    "    qml.CNOT(wires=[2, 0])\n",
+    "    qml.CNOT(wires=[0, 1])\n",
     "    qml.RY(-t1, wires=1)\n",
-    "    qml.CNOT(wires=[0,1])\n",
-    "    qml.CNOT(wires=[1,0])\n",
+    "    qml.CNOT(wires=[0, 1])\n",
+    "    qml.CNOT(wires=[1, 0])\n",
     "\n",
     "# And building our Hamiltonian for the N=3 case\n",
-    "\n",
-    "coeffs = [15.531709 ,-2.1433, -2.1433, 0.21829, -6.125, -9.625, -3.91, -3.91]\n",
-    "obs = [qml.Identity(0),qml.PauliX(0) @ qml.PauliX(1), qml.PauliY(0) @ qml.PauliY(1),qml.PauliZ(0),qml.PauliZ(1),qml.PauliZ(2),qml.PauliX(1) @ qml.PauliX(2),qml.PauliY(1) @ qml.PauliY(2)]\n",
+    "coeffs = [15.531709, -2.1433, -2.1433, 0.21829, -6.125, -9.625, -3.91, -3.91]\n",
+    "obs = [\n",
+    "    qml.Identity(0),\n",
+    "    qml.PauliX(0) @ qml.PauliX(1),\n",
+    "    qml.PauliY(0) @ qml.PauliY(1),\n",
+    "    qml.PauliZ(0),\n",
+    "    qml.PauliZ(1),\n",
+    "    qml.PauliZ(2),\n",
+    "    qml.PauliX(1) @ qml.PauliX(2),\n",
+    "    qml.PauliY(1) @ qml.PauliY(2),\n",
+    "]\n",
     "\n",
     "H = qml.Hamiltonian(coeffs, obs)\n",
-    "cost_fn = qml.ExpvalCost(circuit, H, dev)\n",
     "\n",
+    "#  Define the QNode *directly* with autograd interface\n",
+    "@qml.qnode(dev, interface=\"autograd\")  # Important: Use autograd for QNG\n",
+    "def cost_fn(params):\n",
+    "    circuit(params, wires=dev.wires)\n",
+    "    return qml.expval(H)\n",
+    "    \n",
     "print(H)\n",
     "\n",
     "@qml.qnode(dev)\n",
     "def draw_circuit(params):\n",
-    "    circuit(params,wires=dev.wires)\n",
+    "    circuit(params, wires=dev.wires)\n",
     "    return qml.expval(H)\n",
     "\n",
     "qml.drawer.use_style('default')\n",
-    "print(qml.draw_mpl(draw_circuit)(init_params))"
+    "\n",
+    "# You need to define 'init_params' for the drawing function to work\n",
+    "init_params = np.array([0.1, 0.2])  # Example initial parameters\n",
+    "fig, ax = qml.draw_mpl(draw_circuit)(init_params)\n",
+    "plt.show()"
    ]
   },
   {
@@ -104,7 +120,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -128,7 +144,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -193,7 +209,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -211,22 +227,20 @@
    "source": [
     "opt = qml.QNGOptimizer(stepsize=step_size)\n",
     "\n",
-    "params = init_params\n",
-    "\n",
+    "params = init_params  # Use the same init_params you intend to optimize\n",
     "qng_param_history = [params]\n",
     "qng_cost_history = []\n",
     "\n",
     "for n in range(max_iterations):\n",
-    "\n",
     "    # Take a step in parameter space and record your energy\n",
     "    params, prev_energy = opt.step_and_cost(cost_fn, params)\n",
     "\n",
-    "    # This keeps track of our energy for plotting at comparisons\n",
+    "    # This keeps track of our energy for plotting\n",
     "    qng_param_history.append(params)\n",
     "    qng_cost_history.append(prev_energy)\n",
     "\n",
     "    # Here we see what the energy of our system is with the new parameters\n",
-    "    energy = cost_fn(params)\n",
+    "    energy = cost_fn(params)  # Call the QNode directly\n",
     "\n",
     "    # Calculate difference between new and old energies\n",
     "    conv = np.abs(energy - prev_energy)\n",
@@ -254,7 +268,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {