From ea7ad92a86271803eae23048dec45410a8af4236 Mon Sep 17 00:00:00 2001 From: Abhinav Goel Date: Wed, 2 Oct 2024 13:46:10 -0700 Subject: [PATCH] fixes for mypy --- jax/tools/pgo_nsys_converter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jax/tools/pgo_nsys_converter.py b/jax/tools/pgo_nsys_converter.py index 69f085d3f..10209c9a8 100644 --- a/jax/tools/pgo_nsys_converter.py +++ b/jax/tools/pgo_nsys_converter.py @@ -54,7 +54,7 @@ if __name__ == '__main__': proc.wait() thunk_re = re.compile("hlo_op=(.*)#") - cost_dictionary = dict() + cost_dictionary: dict[str, list] = dict() with open(f"{args.pgle_output_path}", 'w', newline='') as protofile: with open(f"{pgle_folder}{pgle_filename}.pbtxt_{report_name}.csv", newline='') as csvfile: reader = csv.DictReader(csvfile)