Update README.md

README.md

@@ -100,7 +100,7 @@ tokenized_input = tokenizer(input, add_special_tokens= False)
 
 Using model with Jax and Pytorch
 ```python
-from transformers import AutoTokenizer, FlaxAutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, FlaxAutoModelForSequenceClassification
 
 #Load jax model
 model = FlaxAutoModelForSequenceClassification.from_pretrained("Fsoft-AIC/Codebert-docstring-inconsistency")