thijn/sa_ai_training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update train.py

Browse Source
This commit is contained in:
thijn 2025-04-01 09:41:06 +02:00
parent 07e8ef7951
commit f2959f61ae
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@ -82,7 +82,11 @@ for batch in train_dataloader:
batch = {key: value.to(device) for key, value in batch.items()} # Move batch to GPU batch = {key: value.to(device) for key, value in batch.items()} # Move batch to GPU
# Train Model # Train Model
import os
if (os.path.isdir('./results')):
trainer.train(resume_from_checkpoint=True) trainer.train(resume_from_checkpoint=True)
else:
trainer.train()
# Save Model # Save Model
model.save_pretrained('./trained_model') model.save_pretrained('./trained_model')