feat(//cpp/benchmark): Increased workspace size for benchmark, may help

results and increase compatibility Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
pytorch · May 23, 2020 · 8171f79 · 8171f79
1 parent f9af574
commit 8171f79
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/cpp/benchmark/main.cpp b/cpp/benchmark/main.cpp
@@ -136,6 +136,7 @@ int main(int argc, const char* argv[]) {
 
 #ifdef TRT
     auto extra_info = trtorch::ExtraInfo(dims);
+    extra_info.workspace_size = 1 << 24;
 
 #ifdef HALF
     extra_info.op_precision = at::kHalf;