diff --git a/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb b/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
index adeeb5e..c8bf47f 100644
--- a/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
+++ b/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
@@ -1292,14 +1292,24 @@
     }
    ],
    "source": [
-    "from transformers import Seq2SeqTrainer\n",
+    "from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl\n",
+    "from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR\n",
     "\n",
-    "# Save model to local\n",
-    "class PeftSavingCallback(TrainerCallback):\n",
-    "    def on_train_end(self, args, state, control, **kwargs):\n",
-    "        kwargs[\"model\"].save_pretrained(state.best_model_checkpoint)\n",
-    "        pytorch_model_path = os.path.join(state.best_model_checkpoint, \"pytorch_model.bin\")\n",
-    "        os.remove(pytorch_model_path) if os.path.exists(pytorch_model_path) else None\n",
+    "class SavePeftModelCallback(TrainerCallback):\n",
+    "    def on_save(\n",
+    "        self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs,\n",
+    "    ):\n",
+    "        checkpoint_folder = os.path.join(\n",
+    "            args.output_dir, f\"{PREFIX_CHECKPOINT_DIR}-{state.global_step}\"\n",
+    "        )       \n",
+    "\n",
+    "        peft_model_path = os.path.join(checkpoint_folder, \"adapter_model\")\n",
+    "        kwargs[\"model\"].save_pretrained(peft_model_path)\n",
+    "\n",
+    "        pytorch_model_path = os.path.join(checkpoint_folder, \"pytorch_model.bin\")\n",
+    "        if os.path.exists(pytorch_model_path):\n",
+    "            os.remove(pytorch_model_path)\n",
+    "        return control\n",
     "\n",
     "trainer = Seq2SeqTrainer(\n",
     "    args=training_args,\n",
@@ -1309,7 +1319,7 @@
     "    data_collator=data_collator,\n",
     "    # compute_metrics=compute_metrics,\n",
     "    tokenizer=processor.feature_extractor,\n",
-    "    callbacks=[PeftSavingCallback]\n",
+    "    callbacks=[SavePeftModelCallback]\n",
     ")\n",
     "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!"
    ]