Update README.md

README.md

@@ -49,10 +49,12 @@ model_id = "jmtzt/ijepa_vith16_1k"
 processor = AutoProcessor.from_pretrained(model_id)
 model = AutoModel.from_pretrained(model_id)
 
+
 def infer(image):
     inputs = processor(image, return_tensors="pt")
     outputs = model(**inputs)
-    return outputs.pooler_output
+    return outputs.last_hidden_state.mean(dim=1)
+
 
 embed_1 = infer(image_1)
 embed_2 = infer(image_2)