Skip to content

Commit

Permalink
Merge pull request SWivid#401 from hcsolakoglu/fix-finetune-cli-bnb-o…
Browse files Browse the repository at this point in the history
…ptimizer

Add --bnb_optimizer argument to CLI and pass it to Trainer initialization
  • Loading branch information
SWivid authored Nov 5, 2024
2 parents 4a69e6b + dbe35da commit 8a7e849
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/f5_tts/train/finetune_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,14 +55,19 @@ def parse_args():
default=None,
help="Path to custom tokenizer vocab file (only used if tokenizer = 'custom')",
)

parser.add_argument(
"--log_samples",
type=bool,
default=False,
help="Log inferenced samples per ckpt save steps",
)
parser.add_argument("--logger", type=str, default=None, choices=["wandb", "tensorboard"], help="logger")
parser.add_argument(
"--bnb_optimizer",
type=bool,
default=False,
help="Use 8-bit Adam optimizer from bitsandbytes",
)

return parser.parse_args()

Expand Down Expand Up @@ -147,6 +152,7 @@ def main():
wandb_resume_id=wandb_resume_id,
log_samples=args.log_samples,
last_per_steps=args.last_per_steps,
bnb_optimizer=args.bnb_optimizer,
)

train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
Expand Down

0 comments on commit 8a7e849

Please sign in to comment.