fix example

README.md

@@ -76,8 +76,7 @@ Clinical Mosaic was pre-trained on deidentified clinical notes from MIMIC-IV-NOT
 Install the Hugging Face Transformers library and load the model as follows:
 
 ```python
-from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification, BertTokenizer, BertConfig
-import torch
+from transformers import AutoModelForSequenceClassification, BertTokenizer, BertConfig
 
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') # MosaicBERT uses the standard BERT tokenizer
 config = BertConfig.from_pretrained('Sifal/ClinicalMosaic') # the config needs to be passed in
@@ -86,14 +85,19 @@ config = BertConfig.from_pretrained('Sifal/ClinicalMosaic') # the config needs t
 config.num_labels = 4
 config.hidden_size = 768
 
-ClassifierClincalMosaic = AutoModelForSequenceClassification.from_pretrained('Sifal/ClinicalMosaic', config=config, trust_remote_code=True)
+ClassifierClincalMosaic = AutoModelForSequenceClassification.from_pretrained(
+    'Sifal/ClinicalMosaic',
+    config=config,
+    torch_dtype='auto',
+    trust_remote_code=True,
+    device_map="auto"
+)
 
 # Example usage
-
 clinical_text = "..."
 
 inputs = tokenizer(clinical_text, return_tensors="pt")
-embeddings = ClassifierClincalMosaic(**inputs)
+logits = ClassifierClincalMosaic(**inputs).logits
 ```
 
 Further instructions and example scripts are provided in the model’s repository.