mypy on test file

microsoft · Dec 13, 2023 · 4b3f46d · 4b3f46d
1 parent 44146eb
commit 4b3f46d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/datamodules/test_digital_typhoon.py b/tests/datamodules/test_digital_typhoon.py
@@ -37,12 +37,12 @@ def test_split_dataset(self, split_by: str) -> None:
         if split_by == "time":
 
             def find_max_time_per_id(
-                split_sequences: list[dict[str, Union[int, list[int]]]]
-            ) -> dict[int, int]:
+                split_sequences: list[dict[str, Union[str, list[int]]]]
+            ) -> dict[str, int]:
                 # Find the maximum value of each id in train_sequences
-                max_values = {}
+                max_values: dict[str, int] = {}
                 for seq in split_sequences:
-                    id: int = seq["id"]
+                    id: str = str(seq["id"])
                     value: int = max(seq["seq_id"])
                     if id not in max_values or value > max_values[id]:
                         max_values[id] = value