Pytype fix in WB2/compute_statistical_moments.py

PiperOrigin-RevId: 587787044
google-research · Dec 5, 2023 · ea2515b · ea2515b
1 parent 8b40fb6
commit ea2515b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scripts/compute_statistical_moments.py b/scripts/compute_statistical_moments.py
@@ -154,7 +154,8 @@ def main(argv: list[str]) -> None:
       )
       # Rechunk in time
       pcoll_time = pcoll_tmp | f'RechunkTime_{order}' >> xbeam.Rechunk(
-          space_reduce_template.sizes,
+          # Convert to string to satisfy pytype.
+          {str(k): v for k, v in space_reduce_template.sizes.items()},
           reduce_working_chunks,
           time_working_chunks,
           itemsize=RECHUNK_ITEMSIZE.value,