diff --git a/Fine_tune_Llama_2_in_Google_Colab.ipynb b/Fine_tune_Llama_2_in_Google_Colab.ipynb
index 2bafbca..e65d5ec 100644
--- a/Fine_tune_Llama_2_in_Google_Colab.ipynb
+++ b/Fine_tune_Llama_2_in_Google_Colab.ipynb
@@ -6,7 +6,7 @@
       "provenance": [],
       "machine_shape": "hm",
       "gpuType": "V100",
-      "authorship_tag": "ABX9TyMgfvtuquE8AUCpv0te8LOT",
+      "authorship_tag": "ABX9TyPHtqq96zm8/DDNC9+543fd",
       "include_colab_link": true
     },
     "kernelspec": {
@@ -35,7 +35,7 @@
         "# Fine-tune Llama 2 in Google Colab\n",
         "> 🗣️ Large Language Model Course\n",
         "\n",
-        "❤️ Created by [@maximelabonne](), based on Pclanglais' [GitHub Gist](https://gist.github.com/Pclanglais/e90381ed142ee80c8e7ea602b18d50f0).\n"
+        "❤️ Created by [@maximelabonne](), based on Younes Belkada's [GitHub Gist](https://gist.github.com/younesbelkada/9f7f75c94bdc1981c8ca5cc937d4a4da).\n"
       ],
       "metadata": {
         "id": "OSHlAbqzDFDq"
@@ -49,7 +49,7 @@
       },
       "outputs": [],
       "source": [
-        "!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.30.2 trl==0.4.7"
+        "!pip install -q accelerate==0.21.0 peft==0.4.0 bitsandbytes==0.40.2 transformers==4.31.0 trl==0.4.7"
       ]
     },
     {
@@ -189,6 +189,7 @@
         "    device_map=device_map\n",
         ")\n",
         "model.config.use_cache = False\n",
+        "model.config.pretraining_tp = 1\n",
         "\n",
         "peft_config = LoraConfig(\n",
         "    lora_alpha=lora_alpha,\n",
@@ -200,6 +201,8 @@
         "\n",
         "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
         "tokenizer.pad_token = tokenizer.eos_token\n",
+        "# Fix weird overflow issue with fp16 training\n",
+        "tokenizer.padding_side = \"right\"\n",
         "\n",
         "training_arguments = TrainingArguments(\n",
         "    output_dir=output_dir,\n",