From 26f043bd4dd09c985f8852770fa4b3762f56ae6d Mon Sep 17 00:00:00 2001
From: Pablo Montalvo <39954772+molbap@users.noreply.github.com>
Date: Mon, 26 Aug 2024 17:49:44 +0200
Subject: [PATCH] quickfix documentation (#32566)

* fix documentation

* update config
---
 docs/source/en/model_doc/mamba2.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/source/en/model_doc/mamba2.md b/docs/source/en/model_doc/mamba2.md
index edec4872e91900..5ed27881cf18ae 100644
--- a/docs/source/en/model_doc/mamba2.md
+++ b/docs/source/en/model_doc/mamba2.md
@@ -39,11 +39,11 @@ The original code can be found [here](https://github.com/state-spaces/mamba).
 
 ### A simple generation example: 
 ```python 
-from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
+from transformers import Mamba2Config, Mamba2ForCausalLM, AutoTokenizer
 import torch
 model_id = 'mistralai/Mamba-Codestral-7B-v0.1'
 tokenizer = AutoTokenizer.from_pretrained(model_id, revision='refs/pr/9', from_slow=True, legacy=False)
-model = MambaForCausalLM.from_pretrained(model_id, revision='refs/pr/9')
+model = Mamba2ForCausalLM.from_pretrained(model_id, revision='refs/pr/9')
 input_ids = tokenizer("Hey how are you doing?", return_tensors= "pt")["input_ids"]
 
 out = model.generate(input_ids, max_new_tokens=10)