Update README.md

README.md

@@ -101,16 +101,22 @@ inp = jnp.array(inp)
 </details>
 
 
-## Pre-trained checkpoints
+## Usage
+
+The pre-trained models can be used via [Hugging Face hub](https://huggingface.co/collections/apple/aim-65aa3ce948c718a574f09eb7) as follows:
 
-The pre-trained models can be accessed via [PyTorch Hub](https://pytorch.org/hub/) as:
 ```python
-import torch
+from PIL import Image
+
+from aim.torch.models import AIMForImageClassification
+from aim.torch.data import val_transforms
 
-aim_600m = torch.hub.load("apple/ml-aim", "aim_600M")
-aim_1b   = torch.hub.load("apple/ml-aim", "aim_1B")
-aim_3b   = torch.hub.load("apple/ml-aim", "aim_3B")
-aim_7b   = torch.hub.load("apple/ml-aim", "aim_7B")
+img = Image.open(...)
+model = AIMForImageClassification.from_pretrained("apple/aim-7B")
+transform = val_transforms()
+
+inp = transform(img).unsqueeze(0)
+logits, features = model(inp)
 ```
 
 ### Pre-trained backbones