Update README.md

README.md

@@ -101,22 +101,29 @@ inp = jnp.array(inp)
 </details>
 
 
-## Usage
+## Pre-trained checkpoints
 
-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 either via [Hugging Face](https://huggingface.co/collections/apple/aim-65aa3ce948c718a574f09eb7):
 
 ```python
-from PIL import Image
-
+# after running pip install git+https://[email protected]/apple/ml-aim.git
 from aim.torch.models import AIMForImageClassification
-from aim.torch.data import val_transforms
 
-img = Image.open(...)
-model = AIMForImageClassification.from_pretrained("apple/aim-7B")
-transform = val_transforms()
+aim_600m = AIMForImageClassification.from_pretrained("apple/aim-600M")
+aim_1b   = AIMForImageClassification.from_pretrained("apple/aim-1B")
+aim_3b   = AIMForImageClassification.from_pretrained("apple/aim-3B")
+aim_7b   = AIMForImageClassification.from_pretrained("apple/aim-7B")
+```
 
-inp = transform(img).unsqueeze(0)
-logits, features = model(inp)
+or [PyTorch Hub](https://pytorch.org/hub/) as:
+
+```python
+import torch
+
+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")
 ```
 
 ### Pre-trained backbones