Zhen Ye commited on
Commit
d52f7cc
·
1 Parent(s): f78d96f

Use DETR no_timm revision to avoid timm dependency

Browse files
Files changed (1) hide show
  1. models/detectors/detr.py +2 -2
models/detectors/detr.py CHANGED
@@ -21,8 +21,8 @@ class DetrDetector(ObjectDetector):
21
  else:
22
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23
  logging.info("Loading %s onto %s", self.MODEL_NAME, self.device)
24
- self.processor = DetrImageProcessor.from_pretrained(self.MODEL_NAME)
25
- self.model = DetrForObjectDetection.from_pretrained(self.MODEL_NAME)
26
  self.model.to(self.device)
27
  self.model.eval()
28
 
 
21
  else:
22
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23
  logging.info("Loading %s onto %s", self.MODEL_NAME, self.device)
24
+ self.processor = DetrImageProcessor.from_pretrained(self.MODEL_NAME, revision="no_timm")
25
+ self.model = DetrForObjectDetection.from_pretrained(self.MODEL_NAME, revision="no_timm")
26
  self.model.to(self.device)
27
  self.model.eval()
28