From aca175f133dce46e975846523539450969f9d945 Mon Sep 17 00:00:00 2001
From: Naren Dasan <naren@narendasan.com>
Date: Thu, 11 Nov 2021 17:57:51 -0800
Subject: [PATCH] fix(//cpp/bin/torchtrtc): Fix mbs

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>
---
 cpp/bin/torchtrtc/main.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/cpp/bin/torchtrtc/main.cpp b/cpp/bin/torchtrtc/main.cpp
index 27ea60a5a0..80a46837e4 100644
--- a/cpp/bin/torchtrtc/main.cpp
+++ b/cpp/bin/torchtrtc/main.cpp
@@ -288,9 +288,9 @@ int main(int argc, char** argv) {
       "(Repeatable) Module that should always be run in Pytorch for execution (partial compilation must be enabled)",
       {"tem", "torch-executed-mods"});
 
-  args::ValueFlagList<std::string> min_block_size(
+  args::ValueFlag<uint64_t> min_block_size(
       parser,
-      "torch-executed-mods",
+      "min-block-size",
       "Minimum number of contiguous TensorRT supported ops to compile a subgraph to TensorRT",
       {"mbs", "min-block-size"});
 
@@ -626,6 +626,7 @@ int main(int argc, char** argv) {
     std::ofstream out(real_output_path);
     out << engine;
     out.close();
+    return;
   } else {
     auto trt_mod = torchtrt::ts::compile(mod, compile_settings);