Fixup json load

nod-ai · Oct 17, 2024 · dd5dd78 · dd5dd78
1 parent 158a672
commit dd5dd78
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/models/turbine_models/custom_models/torchbench/export.py b/models/turbine_models/custom_models/torchbench/export.py
@@ -420,7 +420,7 @@ def run_main(model_id, args, tb_dir, tb_args):
     if args.compile_to in ["torch", "mlir"]:
         safe_name = utils.create_safe_name(
             model_id,
-            f"_{static_dim}_{precision}",
+            f"_{static_dim}_{args.precision}",
         )
         with open(f"{safe_name}.mlir", "w+") as f:
             f.write(mod_str)
@@ -445,11 +445,9 @@ def run_main(model_id, args, tb_dir, tb_args):
     from turbine_models.custom_models.torchbench.cmd_opts import args, unknown
     import json
 
-    torchbench_models_dict = json.load(args.model_list_json)
     for list in args.model_lists:
-        torchbench_models_dict = json.load(list)
-        with open(args.models_json, "r") as f:
-            torchbench_models_dict = json.load(file)
+        with open(list, "r") as f:
+            torchbench_models_dict = json.load(f)
 
         tb_dir = setup_torchbench_cwd()
         if args.model_id.lower() == "all":