Skip to content

Commit 57909c1

Browse files
fix yolo-v4 demo (#2546)
1 parent da693f8 commit 57909c1

File tree

1 file changed

+3
-4
lines changed

1 file changed

+3
-4
lines changed

demos/common/python/models/yolo.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -245,10 +245,9 @@ def _get_output_info(self):
245245

246246
output_info = {}
247247
for i, (name, layer) in enumerate(outputs):
248-
if name.startswith('conv'):
249-
shape = layer.shape
250-
yolo_params = self.Params(num, shape[2:4], anchors, masks[len(output_info)])
251-
output_info[name] = (shape, yolo_params)
248+
shape = layer.shape
249+
yolo_params = self.Params(num, shape[2:4], anchors, masks[len(output_info)])
250+
output_info[name] = (shape, yolo_params)
252251
return output_info
253252

254253
@staticmethod

0 commit comments

Comments
 (0)