fix tritonbench parser

Summary: this doesn't work if `--isolate` is the last arg Reviewed By: xuzhao9 Differential Revision: D65837725 fbshipit-source-id: f24728ddb63623d8beb339feec12b35ca635d491
pytorch-labs · Nov 13, 2024 · f63be70 · f63be70
1 parent eba50ab
commit f63be70
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tritonbench/utils/parser.py b/tritonbench/utils/parser.py
@@ -200,7 +200,7 @@ def _find_param_loc(params, key: str) -> int:
 def _remove_params(params, loc):
     if loc == -1:
         return params
-    if params[loc + 1].startswith("--"):
+    if (loc + 1) < len(params) and params[loc + 1].startswith("--"):
         return params[:loc] + params[loc + 1 :]
     return params[:loc] + params[loc + 2 :]