diff --git a/micro_bench.ipynb b/micro_bench.ipynb
index 6ff62ff..7f702b4 100644
--- a/micro_bench.ipynb
+++ b/micro_bench.ipynb
@@ -305,7 +305,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from transformers import AutoConfig, BertLayer\n",
+    "from transformers import AutoConfig\n",
+    "from transformers.models.bert.modeling_bert import BertLayer\n",
     "\n",
     "config = AutoConfig.from_pretrained(\"bert-large-uncased\")\n",
     "layer = BertLayer(config).half().cuda()"