Fix some issues with benchmark data output (#13641)

Signed-off-by: Huy Do <huydhn@gmail.com>
This commit is contained in:
Huy Do
2025-02-23 18:23:18 -08:00
committed by GitHub
parent cbae7af552
commit e7ef74e26e
7 changed files with 61 additions and 16 deletions

View File

@ -11,7 +11,7 @@ from typing import Any, Dict, List, Optional, Tuple
import torch
import uvloop
from benchmark_utils import convert_to_pytorch_benchmark_format
from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
from PIL import Image
from tqdm import tqdm
from transformers import (AutoModelForCausalLM, AutoTokenizer,
@ -366,8 +366,7 @@ def save_to_pytorch_benchmark_format(args: argparse.Namespace,
if pt_records:
# Don't use json suffix here as we don't want CI to pick it up
pt_file = f"{os.path.splitext(args.output_json)[0]}.pytorch.json"
with open(pt_file, "w") as f:
json.dump(pt_records, f)
write_to_json(pt_file, pt_records)
def main(args: argparse.Namespace):