Merge pull request #47 from mlfoundations/jean/fix_batch_size

Small fix batch_size type
mlfoundations · Jan 9, 2025 · ea417a4 · ea417a4
2 parents e937b74 + 09b66c2
commit ea417a4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/eval/eval.py b/eval/eval.py
@@ -252,10 +252,10 @@ def cli_evaluate(args: Optional[argparse.Namespace] = None) -> None:
         with open(args.config, "r") as file:
             tasks_yaml = yaml.safe_load(file)
         args.tasks = ",".join([t["task_name"] for t in tasks_yaml["tasks"]])
-        batch_sizes_list = [t["batch_size"] for t in tasks_yaml["tasks"]]
+        batch_sizes_list = [int(t["batch_size"]) for t in tasks_yaml["tasks"]]
         args.annotator_model = tasks_yaml.get("annotator_model", args.annotator_model)
     else:
-        batch_sizes_list = [args.batch_size for _ in range(len(args.tasks.split(",")))]
+        batch_sizes_list = [int(args.batch_size) for _ in range(len(args.tasks.split(",")))]
 
     # Initialize evaluation tracker
     if args.output_path: