bug fix

MarcusLoppe · Jun 17, 2024 · 481428d · 481428d
1 parent b7a3611
commit 481428d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/meshgpt_pytorch/trainer.py b/meshgpt_pytorch/trainer.py
@@ -86,7 +86,7 @@ def __init__(
         accelerator_kwargs: dict = dict(),
         optimizer_kwargs: dict = dict(),
         checkpoint_every = 1000,
-        checkpoint_every_epoch: Optional[int] = None,
+        checkpoint_every_epoch: Type[int] | None = None,
         checkpoint_folder = './checkpoints',
         data_kwargs: Tuple[str, ...] = ('vertices', 'faces', 'face_edges'),
         warmup_steps = 1000,
@@ -436,14 +436,14 @@ def __init__(
         val_dataset: Dataset | None = None,
         val_every = 1,
         val_num_batches = 5,
-        scheduler: Optional[Type[_LRScheduler]] = None,
+        scheduler: Type[_LRScheduler] | None = None,
         scheduler_kwargs: dict = dict(),
         ema_kwargs: dict = dict(),
         accelerator_kwargs: dict = dict(),
         optimizer_kwargs: dict = dict(),
 
         checkpoint_every = 1000, 
-        checkpoint_every_epoch: Optional[int] = None,
+        checkpoint_every_epoch: Type[int] | None = None,
         checkpoint_folder = './checkpoints',
         data_kwargs: Tuple[str, ...] = ('vertices', 'faces', 'face_edges', 'text'),
         warmup_steps = 1000,