Update README.md

README.md

@@ -57,15 +57,16 @@ Example usage:
 
 ```
 import cv2
+import torch
 from transformers import AutoModel
 
-def crop_mammo(img, crop_model, device):
-  img_shape = torch.tensor([img.shape[:2]])
-  x = crop_model.preprocess(img)
+def crop_mammo(img, model, device):
+  img_shape = torch.tensor([img.shape[:2]]).to(device)
+  x = model.preprocess(img)
   x = torch.from_numpy(x).expand(1, 1, -1, -1).float().to(device)
   with torch.inference_mode():
     coords = model(x, img_shape)
-  coords = coords[0].numpy()
+  coords = coords[0].cpu().numpy()
   x, y, w, h = coords
   return img[y: y + h, x: x + w]