Minor fixes

zhiqwang · Dec 26, 2021 · 7215fb4 · 7215fb4
1 parent 64c7c9b
commit 7215fb4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/yolort/runtime/trt_helper.py b/yolort/runtime/trt_helper.py
@@ -31,6 +31,9 @@
 logger = logging.getLogger("TRTHelper")
 
 
+__all__ = ["YOLOTRTModule", "EngineBuilder"]
+
+
 class YOLOTRTModule(nn.Module):
     """
     TensorRT deployment friendly wrapper for YOLO.