nielsr HF staff commited on
Commit
e76a695
1 Parent(s): 33b593c

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -31,7 +31,7 @@ from PIL import Image
31
  from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
32
 
33
  model_id = "IDEA-Research/grounding-dino-base"
34
- device = "cuda:0"
35
 
36
  processor = AutoProcessor.from_pretrained(model_id)
37
  model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device)
 
31
  from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
32
 
33
  model_id = "IDEA-Research/grounding-dino-base"
34
+ device = "cuda" if torch.cuda.is_available() else "cpu"
35
 
36
  processor = AutoProcessor.from_pretrained(model_id)
37
  model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device)