From 00e803f2a230f2d1c3f59e535d63053a01666ae0 Mon Sep 17 00:00:00 2001
From: twanchen <twanchen@google.com>
Date: Wed, 10 Jul 2024 13:10:51 +0800
Subject: [PATCH] Fix Unknown Architecture Error

---
 site/en/gemma/docs/pytorch_gemma.ipynb | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/site/en/gemma/docs/pytorch_gemma.ipynb b/site/en/gemma/docs/pytorch_gemma.ipynb
index 3286096c7..b5b432de9 100644
--- a/site/en/gemma/docs/pytorch_gemma.ipynb
+++ b/site/en/gemma/docs/pytorch_gemma.ipynb
@@ -220,9 +220,11 @@
       "outputs": [],
       "source": [
         "import torch\n",
+        "from gemma import config as gemma_config\n",
         "\n",
         "# Set up model config.\n",
         "model_config = get_config_for_2b() if \"2b\" in VARIANT else get_config_for_7b()\n",
+        "model_config.architecture = gemma_config.Architecture.GEMMA_1\n",
         "model_config.tokenizer = tokenizer_path\n",
         "model_config.quant = 'quant' in VARIANT\n",
         "\n",