Minor fixes

zhiqwang · Jan 24, 2022 · a3f80fb · a3f80fb
1 parent c72dd6a
commit a3f80fb
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/deployment/tensorrt/main.cpp b/deployment/tensorrt/main.cpp
@@ -138,7 +138,7 @@ ICudaEngine* CreateCudaEngineFromOnnx(
   if (!config) {
     return nullptr;
   }
-  config->setMaxWorkspaceSize(20 * (1U << 20));
+  config->setMaxWorkspaceSize(40 * (1U << 20));
   config->setFlag(BuilderFlag::kGPU_FALLBACK);
   if (enable_int8) {
     if (builder->platformHasFastInt8()) {
@@ -267,7 +267,7 @@ std::vector<Detection> YOLOv5Detector::detect(cv::Mat& image) {
   int32_t num_detections = 0;
   std::vector<float> detection_boxes;
   std::vector<float> detection_scores;
-  std::vector<float> detection_labels;
+  std::vector<int> detection_labels;
 
   Dims dim = engine->getBindingDimensions(0);
   dim.d[0] = batch_size;

diff --git a/yolort/runtime/yolo_graphsurgeon.py b/yolort/runtime/yolo_graphsurgeon.py
@@ -140,7 +140,7 @@ def register_nms(
         Register the ``EfficientNMS_TRT`` plugin node.
 
         NMS expects these shapes for its input tensors:
-            - box_net: [batch_size, number_boxes, 1, 4]
+            - box_net: [batch_size, number_boxes, 4]
             - class_net: [batch_size, number_boxes, number_labels]
 
         Args: