diff --git a/docs/tutorials/tfr_bert.ipynb b/docs/tutorials/tfr_bert.ipynb
index 804a368..d37a1a6 100644
--- a/docs/tutorials/tfr_bert.ipynb
+++ b/docs/tutorials/tfr_bert.ipynb
@@ -346,7 +346,7 @@
         "vali_dataset = task.build_inputs(task_config.validation_data)\n",
         "train_iterator = iter(train_dataset)\n",
         "vali_iterator = iter(vali_dataset)\n",
-        "optimizer = tf.keras.optimizers.Adam(lr=1e-6)\n",
+        "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-6)\n",
         "\n",
         "NUM_TRAIN_STEPS = 100\n",
         "EVAL_STEPS = 10\n",