trl-4-dnd / quickstart.py
vishaljoshi24's picture
Amended quickstart.py for better memory allocation
4a4b568
raw
history blame contribute delete
293 Bytes
import os
from trl import SFTTrainer
from datasets import load_dataset
dataset = load_dataset("trl-lib/Capybara", split="train")
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
trainer = SFTTrainer(
model="Qwen/Qwen3-0.6B",
train_dataset=dataset,
)
trainer.train()