Skip to content

Commit aca175f

Browse files
committed
fix(//cpp/bin/torchtrtc): Fix mbs
Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
1 parent 9eae269 commit aca175f

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

cpp/bin/torchtrtc/main.cpp

+3-2
Original file line numberDiff line numberDiff line change
@@ -288,9 +288,9 @@ int main(int argc, char** argv) {
288288
"(Repeatable) Module that should always be run in Pytorch for execution (partial compilation must be enabled)",
289289
{"tem", "torch-executed-mods"});
290290

291-
args::ValueFlagList<std::string> min_block_size(
291+
args::ValueFlag<uint64_t> min_block_size(
292292
parser,
293-
"torch-executed-mods",
293+
"min-block-size",
294294
"Minimum number of contiguous TensorRT supported ops to compile a subgraph to TensorRT",
295295
{"mbs", "min-block-size"});
296296

@@ -626,6 +626,7 @@ int main(int argc, char** argv) {
626626
std::ofstream out(real_output_path);
627627
out << engine;
628628
out.close();
629+
return;
629630
} else {
630631
auto trt_mod = torchtrt::ts::compile(mod, compile_settings);
631632

0 commit comments

Comments
 (0)