Merge pull request #3 from rballachay/patch-1

should pass device in case of cpu
luca-medeiros · Apr 14, 2023 · 83c6eaf · 83c6eaf
2 parents 5bfbfad + d0149b2
commit 83c6eaf
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lang_sam/lang_sam.py b/lang_sam/lang_sam.py
@@ -75,7 +75,8 @@ def predict_dino(self, image_pil, text_prompt, box_threshold, text_threshold):
  image=image_trans,
  caption=text_prompt,
  box_threshold=box_threshold,
- text_threshold=text_threshold)
+ text_threshold=text_threshold,
+ device=self.device)
  W, H = image_pil.size
  boxes = box_ops.box_cxcywh_to_xyxy(boxes) * torch.Tensor([W, H, W, H])