Updated README.md

README.md

@@ -17,6 +17,6 @@ pruned-<original model name>-<pruning method>-<sparsity ratio>-<sparsity type>-<
 Load the model using the `transformers` library, which can be installed using `pip install transformers`:
 ```python
 # main.py
-from transformers import AutoModel
-model = AutoModel.from_pretrained("multilingual-pruning/<model name>")
+from transformers import AutoModelForCausalLM
+model = AutoModelForCausalLM.from_pretrained("multilingual-pruning/<model name>")
 ```