Fix typing annotation

nucleons · Oct 16, 2023 · 7c10aad · 7c10aad
1 parent c1c618c
commit 7c10aad
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/pretrain/openwebtext.py b/pretrain/openwebtext.py
@@ -69,7 +69,7 @@ def setup(devices: int = 1, precision: Optional[str] = None, resume: Union[bool,
     fabric.launch(main, resume=resume)
 
 
-def main(fabric: L.Fabric, resume: bool) -> None:
+def main(fabric: L.Fabric, resume: Union[bool, Path]) -> None:
     speed_monitor = SpeedMonitor(fabric, window_size=50, time_unit="seconds")
 
     if fabric.global_rank == 0:

diff --git a/pretrain/redpajama.py b/pretrain/redpajama.py
@@ -86,7 +86,7 @@ def setup(
     fabric.launch(main, train_data_dir, val_data_dir, resume)
 
 
-def main(fabric: L.Fabric, train_data_dir: Path, val_data_dir: Path, resume: bool) -> None:
+def main(fabric: L.Fabric, train_data_dir: Path, val_data_dir: Path, resume: Union[bool, Path]) -> None:
     speed_monitor = SpeedMonitor(fabric, window_size=50, time_unit="seconds")
 
     if fabric.global_rank == 0:

diff --git a/setup.py b/setup.py
@@ -15,7 +15,7 @@
     url="https://github.com/lightning-AI/lit-gpt",
     install_requires=[
         "torch>=2.1.0",
-        "lightning @ git+https://github.com/Lightning-AI/lightning@71aed751f7f0ca8422ddca256e602099070f490b"
+        "lightning @ git+https://github.com/Lightning-AI/lightning@71aed751f7f0ca8422ddca256e602099070f490b",
     ],
     packages=find_packages(),
     long_description=readme,