From f63be702d041c5471a4814a6f9e2250cc4484877 Mon Sep 17 00:00:00 2001
From: Nicolas Macchioni <nmacchioni@meta.com>
Date: Wed, 13 Nov 2024 09:03:50 -0800
Subject: [PATCH] fix tritonbench parser

Summary: this doesn't work if `--isolate` is the last arg

Reviewed By: xuzhao9

Differential Revision: D65837725

fbshipit-source-id: f24728ddb63623d8beb339feec12b35ca635d491
---
 tritonbench/utils/parser.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tritonbench/utils/parser.py b/tritonbench/utils/parser.py
index 801e8ccf..d5389c3f 100644
--- 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 :]