(self, ckpt_path: str)
| 492 | python_ckpt_map: dict[str, Any] |
| 493 | |
| 494 | def _spec_from_path(self, ckpt_path: str): |
| 495 | # TODO(markblee): Enable ocdbt driver. |
| 496 | return array_serialization.get_tensorstore_spec(ckpt_path) |
| 497 | |
| 498 | def _get_spec(self, step: int, state: NestedTensor, ckpt_dir: str) -> CheckpointSpec: |
| 499 | spec = self.CheckpointSpec( |