Fix disentangling notebook typo

PiperOrigin-RevId: 346187900
google · Dec 13, 2020 · 2a8e5f5 · 2a8e5f5
1 parent f5ab98f
commit 2a8e5f5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/notebooks/Disentangling_Trainability_and_Generalization.ipynb b/notebooks/Disentangling_Trainability_and_Generalization.ipynb
@@ -1465,7 +1465,7 @@
         "    lambda_max = np.linalg.eigh(kdd.ntk)[0][-1]\n",
         "    eta = y_train.size * 2. / lambda_max\n",
         "    predict_fn = nt.predict.gradient_descent_mse(kdd.ntk, y_train)\n",
-        "    predict_fn = partial(predict_fn, fx_test_0=0., ntk_test_train=ktd.ntk)\n",
+        "    predict_fn = partial(predict_fn, fx_test_0=0., k_test_train=ktd.ntk)\n",
         "    train, test = vmap(predict_fn)(ts)\n",
         "\n",
         "    # Compute the next kernel after iterating the map for dlayers.\n",