Update README.md

README.md

@@ -48,11 +48,12 @@ This repository contains the **MobileCLIP2-S2** checkpoint.
 import torch
 import open_clip
 from PIL import Image
+from urllib.request import urlopen
 from timm.utils import reparameterize_model
 
-model, _, preprocess = open_clip.create_model_and_transforms('MobileCLIP2-S2', pretrained='dfndr2b')
+model, _, preprocess = open_clip.create_model_and_transforms('MobileCLIP2-S0', pretrained='dfndr2b')
 model.eval()
-tokenizer = open_clip.get_tokenizer('MobileCLIP2-S2')
+tokenizer = open_clip.get_tokenizer('MobileCLIP2-S0')
 
 # For inference/model exporting purposes, optionally reparameterize for better performance
 model = reparameterize_model(model)
@@ -63,7 +64,7 @@ image = Image.open(urlopen(
 image = preprocess(image).unsqueeze(0)
 text = tokenizer(["a diagram", "a dog", "a cat", "a doughnut"])
 
-with torch.no_grad(), torch.cuda.amp.autocast():
+with torch.no_grad(), torch.amp.autocast(image.device.type):
     image_features = model.encode_image(image)
     text_features = model.encode_text(text)
     image_features /= image_features.norm(dim=-1, keepdim=True)