minor

RonSheely · Jul 25, 2022 · da5a16f · da5a16f
1 parent 2c1df24
commit da5a16f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mcunet/utils/det_helper.py b/mcunet/utils/det_helper.py
@@ -2,7 +2,9 @@
 import torch.nn as nn
 import numpy as np
 
-__all__ = ['standard_nms', 'StandardNMS', 'MergeNMS']
+# Some of the code is adapted from gluoncv: https://cv.gluon.ai/model_zoo/detection.html
+
+__all__ = ['standard_nms', 'StandardNMS', 'MergeNMS', 'Yolo3Output']
 
 
 def bbox_iou(box1: torch.Tensor, box2: torch.Tensor, offset=0) -> torch.Tensor: