storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-24 08:25:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support checkpointing.

Browse Source
This commit is contained in:
sayakpaul
2024-12-01 20:34:35 +05:30
parent c8a118cefd
commit 79b546b264
2 changed files with 55 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/mochi-1/args.py
+11 -6
View File
@@ -197,6 +197,16 @@ def _get_training_args(parser: argparse.ArgumentParser) -> None:
default=200,
help="Number of steps for the warmup in the lr scheduler.",
)
parser.add_argument(
"--checkpointing_steps",
type=int,
default=None,
)
parser.add_argument(
"--resume_from_checkpoint",
type=str,
default=None,
)
def _get_optimizer_args(parser: argparse.ArgumentParser) -> None:
@@ -242,12 +252,7 @@ def _get_configuration_args(parser: argparse.ArgumentParser) -> None:
" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
),
)
parser.add_argument(
"--report_to",
type=str,
default=None,
help="If logging to wandb."
)
parser.add_argument("--report_to", type=str, default=None, help="If logging to wandb.")
def get_args():