This commit is contained in:
Glenn Jocher 2019-08-06 16:57:33 +02:00
parent 082fdebfc1
commit 141032045b
1 changed files with 1 additions and 1 deletions

View File

@ -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