Update README.md

README.md

@@ -60,10 +60,10 @@ The model can be used for text generation via its integrated `generate()` method
 ### Example: Direct Inference
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
 
 # Load the model and tokenizer from the hub
-model = LatentRecurrentDepthModel.from_pretrained("codewithdark/latent-recurrent-depth-lm")
+model = AutoModelForCausalLM.from_pretrained("codewithdark/latent-recurrent-depth-lm")
 tokenizer = AutoTokenizer.from_pretrained("codewithdark/latent-recurrent-depth-lm")
 
 prompt = "In the realm of language modeling"
@@ -87,7 +87,7 @@ print(generated_text)
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained("codewithdark/latent-recurrent-depth-lm")
-model = LatentRecurrentDepthModel.from_pretrained("codewithdark/latent-recurrent-depth-lm")
+model = AutoModelForCausalLM.from_pretrained("codewithdark/latent-recurrent-depth-lm")
 
 prompt = "In the realm of language modeling"
 input_ids = tokenizer(prompt, return_tensors="pt").input_ids