Update README.md

README.md

@@ -49,6 +49,7 @@ This repository contains the **MobileCLIP2-S0** 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-S0', pretrained='dfndr2b')
@@ -64,7 +65,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)