Update README.md

README.md

@@ -111,8 +111,12 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
 text = "الدنيا مش مستاهلة تجري كده، خد وقتك واستمتع بالحاجة البسيطة"
 inputs = tokenizer(text, return_tensors="pt")
-outputs = model(**inputs)
-pred = torch.argmax(outputs.logits, dim=1).item()
+
+# Run inference
+with torch.inference_mode():
+    logits = model(**inputs).logits
+
+pred = torch.argmax(logits, dim=-1).item()
 
 print(f"Predicted Dialect: {model.config.id2label[pred]}")
 ```