henry000 commited on
Commit
95d2860
·
1 Parent(s): a421f81

👽️ [Fix] auto anchor device bug, need improve

Browse files
Files changed (1) hide show
  1. yolo/utils/bounding_box_utils.py +2 -1
yolo/utils/bounding_box_utils.py CHANGED
@@ -314,7 +314,8 @@ class Vec2Box:
314
 
315
  def create_auto_anchor(self, model: YOLO, image_size):
316
  W, H = image_size
317
- dummy_input = torch.zeros(1, 3, H, W).to(self.device)
 
318
  dummy_output = model(dummy_input)
319
  strides = []
320
  for predict_head in dummy_output["Main"]:
 
314
 
315
  def create_auto_anchor(self, model: YOLO, image_size):
316
  W, H = image_size
317
+ # TODO: need accelerate dummy test
318
+ dummy_input = torch.zeros(1, 3, H, W)
319
  dummy_output = model(dummy_input)
320
  strides = []
321
  for predict_head in dummy_output["Main"]: