Update README.md

README.md

@@ -28,15 +28,16 @@ You can use the model using pipeline or AutoModelForMultimodalLM options offered
 # Use a pipeline as a high-level helper
 from transformers import pipeline
 
-pipe = pipeline("text-generation", model="pulkitchowdry/sample-agnews-classifer")
+pipe = pipeline("text-classification", model="pulkitchowdry/sample-agnews-classifer")
 text = "US-Iran sign a peace deal"
 pipe(text)
 
 # Load model directly
-from transformers import AutoTokenizer, AutoModelForMultimodalLM
+from transformers import AutoTokenizer, AutoModel
 
 tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
-model = AutoModelForMultimodalLM.from_pretrained("pulkitchowdry/sample-agnews-classifer")
+model = AutoModel.from_pretrained("pulkitchowdry/sample-agnews-classifer", dtype="auto")
+
 text = "US-Iran sign a peace deal"
 inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
 prediction = model(inputs)