From dcc2e99fb25251a4e7363a814f6e3b038987c48a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 11 Apr 2020 10:55:49 -0700 Subject: [PATCH] get_yolo_layers() --- models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models.py b/models.py index 4ecf99e6..90d94ebe 100755 --- a/models.py +++ b/models.py @@ -334,7 +334,7 @@ class Darknet(nn.Module): def get_yolo_layers(model): - return [i for i, x in enumerate(model.module_defs) if x['type'] == 'yolo'] # [82, 94, 106] for yolov3 + return [i for i, m in enumerate(model.module_list) if m.__class__.__name__ == 'YOLOLayer'] # [89, 101, 1113] def load_darknet_weights(self, weights, cutoff=-1):