Update README.md

README.md

@@ -20,6 +20,12 @@ from PIL import Image
 import torch
 from transformers import  AutoImageProcessor, AutoModelForObjectDetection
 
+device = 'cpu'
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch.backends.mps.is_available():
+    device = torch.device('mps')
+
 ckpt = 'yainage90/fashion-object-detection'
 image_processor = AutoImageProcessor.from_pretrained(ckpt)
 model = AutoModelForObjectDetection.from_pretrained(ckpt).to(device)