Fix import bug in data_utils.py

Summary: As the title says

Reviewed By: nmacchioni

Differential Revision: D65881268

fbshipit-source-id: 91ab130b133e2d35e15244d971882f3b51946331

Author: xuzhao9

tritonbench/operators/softmax/operator.py

@@ -103,7 +103,7 @@ def _inner():
 
     def get_input_iter(self):
         M = 4096
-        shapes = (tuple(M, 128 * i) for i in range(2, 100))
+        shapes = [(M, 128 * i) for i in range(2, 100)]
         if IS_FBCODE and self.tb_args.production_shapes:
             shapes = get_production_shapes(self.name, "softmax")
         for M, N in shapes:

tritonbench/utils/data_utils.py

@@ -1,4 +1,4 @@
-from .triton_ops import IS_FBCODE
+from .triton_op import IS_FBCODE
 
 
 def get_production_shapes(op_name, op_type):

tritonbench/utils/parser.py

@@ -200,8 +200,12 @@ def _find_param_loc(params, key: str) -> int:
 def _remove_params(params, loc):
     if loc == -1:
         return params
-    if (loc + 1) < len(params) and params[loc + 1].startswith("--"):
+    if loc == len(params) - 1:
+        return params[:loc]
+    if params[loc + 1].startswith("--"):
         return params[:loc] + params[loc + 1 :]
+    if loc == len(params) - 2:
+        return params[:loc]
     return params[:loc] + params[loc + 2 :]