From 141032045b53e4499f6ec82f915e3467f6031f3d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 6 Aug 2019 16:57:33 +0200 Subject: [PATCH] updates --- models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models.py b/models.py index 1809b1f3..48a7984a 100755 --- a/models.py +++ b/models.py @@ -203,7 +203,7 @@ class Darknet(nn.Module): return output elif ONNX_EXPORT: output = torch.cat(output, 1) # cat 3 layers 85 x (507, 2028, 8112) to 85 x 10647 - nc = self.module_list[self.yolo_layers[0]][0].nc # number of classes + nc = self.module_list[self.yolo_layers[0]].nc # number of classes return output[5:5 + nc].t(), output[:4].t() # ONNX scores, boxes else: io, p = list(zip(*output)) # inference output, training output