Minor typing cleanup

triton-inference-server · Oct 16, 2024 · 3949897 · 3949897
1 parent 1984d16
commit 3949897
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/genai-perf/genai_perf/checkpoint/checkpoint.py b/genai-perf/genai_perf/checkpoint/checkpoint.py
@@ -19,6 +19,7 @@
 from genai_perf.config.input.config_command import ConfigCommand
 from genai_perf.config.run.results import Results
 from genai_perf.exceptions import GenAIPerfException
+from genai_perf.types import CheckpointObject
 
 
 @dataclass(frozen=True)
@@ -60,10 +61,12 @@ def _read_from_checkpoint(self) -> None:
             try:
                 with open(checkpoint_filename, "r") as checkpoint_file:
                     checkpoint_json = json.load(checkpoint_file)
-                    self._state = {}
-                    self._state["Results"] = Results.read_from_checkpoint(
-                        checkpoint_json["Results"]
-                    )
+                    self._state: CheckpointObject = {
+                        "Results": Results.read_from_checkpoint(
+                            checkpoint_json["Results"]
+                        )
+                    }
+
             except EOFError:
                 raise (
                     GenAIPerfException(