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):