From 89d249d670e4c37a32b71cc44e4afc85c784efeb Mon Sep 17 00:00:00 2001
From: WolframRhodium <WolframRhodium@users.noreply.github.com>
Date: Fri, 21 Jan 2022 15:35:58 +0800
Subject: [PATCH] vstrt/trt_utils.h: add support for static shapes

---
 vstrt/trt_utils.h | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/vstrt/trt_utils.h b/vstrt/trt_utils.h
index e03de19..b9ea242 100644
--- a/vstrt/trt_utils.h
+++ b/vstrt/trt_utils.h
@@ -238,7 +238,23 @@ std::variant<ErrorMessage, InferenceInstance> getInstance(
         }
         exec_context->setBindingDimensions(0, dims);
     } else if (std::holds_alternative<RequestedTileSize>(tile_size)) {
-        return set_error("Engine has no dynamic dimensions");
+        is_dynamic = false;
+
+        nvinfer1::Dims dims = exec_context->getBindingDimensions(0);
+
+        if (std::holds_alternative<RequestedTileSize>(tile_size)) {
+            if (dims.d[2] != std::get<RequestedTileSize>(tile_size).tile_h ||
+                dims.d[3] != std::get<RequestedTileSize>(tile_size).tile_w
+            ) {
+                return set_error("requested tile size not applicable");
+            }
+        } else {
+            if (dims.d[2] != std::get<VideoSize>(tile_size).height ||
+                dims.d[3] != std::get<VideoSize>(tile_size).width
+            ) {
+                return set_error("not supported video dimensions");
+            }
+        }
     }
 
     MemoryResource src {};