Use device-agnostic code instead of hardcoded .cuda()

README.md

@@ -21,18 +21,20 @@ pip install transformers torch torchvision sentencepiece
 ```
 
 ```python
+import torch
 from transformers import AutoModel
 from torchvision import transforms
 from PIL import Image
 import requests
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModel.from_pretrained("google/tipsv2-g14-dpt", trust_remote_code=True)
-model.eval().cuda()
+model.eval().to(device)
 
 url = "https://raw.githubusercontent.com/google-deepmind/tips/main/scenic/images/example_image_2.jpg"
 image = Image.open(requests.get(url, stream=True).raw)
 transform = transforms.Compose([transforms.Resize((448, 448)), transforms.ToTensor()])
-pixel_values = transform(image).unsqueeze(0).cuda()
+pixel_values = transform(image).unsqueeze(0).to(device)
 
 # All tasks at once
 outputs = model(pixel_values)