File size: 630 Bytes
167d5e2 a96669d 167d5e2 a96669d 167d5e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from transformers import TrainingArguments
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
warmup_steps=500,
weight_decay=0.01,
logging_dir='./logs',
logging_steps=50,
save_steps=500,
eval_steps=500,
evaluation_strategy="steps",
save_strategy="steps",
save_total_limit=3,
load_best_model_at_end=True,
learning_rate=5e-5,
fp16=True,
gradient_checkpointing=True,
# Remove CPU-only settings
no_cuda=False, # Allow GPU usage
use_cpu=False # Allow GPU usage
) |