diff --git a/models/yolo.py b/models/yolo.py index ad3d8fee7..a37be54ac 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -679,12 +679,15 @@ def parse_model(d, ch): # model_dict, input_channels(3) n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain if m in { - Conv, AConv, ConvTranspose, + Conv, AConv, ConvTranspose, CBLinear, Bottleneck, SPP, SPPF, DWConv, BottleneckCSP, nn.ConvTranspose2d, DWConvTranspose2d, SPPCSPC, ADown, RepNCSPELAN4, SPPELAN}: c1, c2 = ch[f], args[0] if c2 != no: # if not output - c2 = make_divisible(c2 * gw, 8) + if isinstance(c2, list): + c2 = [make_divisible(elem * gw, 8) for elem in c2] + else: + c2 = make_divisible(c2 * gw, 8) args = [c1, c2, *args[1:]] if m in {BottleneckCSP, SPPCSPC}: