AI/gpt4all
1
0
Fork 0
mirror of https://github.com/nomic-ai/gpt4all.git synced 2024-10-01 01:06:10 -04:00
Code Issues Packages Projects Releases Wiki Activity

fix: add epoch train

Browse Source
This commit is contained in:
Zach Nussbaum 2023-03-27 16:32:35 +00:00
parent bb28929305
commit 24765a1965
1 changed files with 41 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
train.py
View File

@ -115,6 +115,7 @@ def train(accelerator, config):
"gradient_accumulation_steps"
]
for epoch in range(config["num_epochs"]):
for step, batch in enumerate(tqdm(train_dataloader)):
model.train()
outputs = model(**batch)
@ -158,6 +159,13 @@ def train(accelerator, config):
train_loss.reset()
accelerator.print(f"Epoch {epoch} finished")
accelerator.print(f"Pushing to HF hub")
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
if accelerator.is_main_process:
unwrapped_model.push_to_hub(config["save_name"], private=True)
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
@ -168,6 +176,7 @@ def train(accelerator, config):
state_dict=accelerator.get_state_dict(model),
)
if accelerator.is_main_process:
unwrapped_model.push_to_hub(config["save_name"], private=True)
accelerator.end_training()