make mypy happy

allenai · Jan 23, 2025 · 94d756b · 94d756b
1 parent 3f097bc
commit 94d756b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/olmo_core/distributed/checkpoint/__init__.py b/src/olmo_core/distributed/checkpoint/__init__.py
@@ -459,6 +459,8 @@ def save(state_dict: Dict[str, Any], path: Path):
 
     def get_chunks(prefix: str) -> Tuple[Path, List[Tuple[Path, List[str]]]]:
         assert unshard_strategy is not None
+        assert isinstance(target_dir, Path)
+
         if unshard_strategy.name == UnshardStrategyType.one_file:
             path = target_dir / f"{prefix}.{ext}"
             return path, [(path, [prefix])]