from ultralytics import YOLO


model = YOLO("yolov8n.pt")
model.to('cuda')
model.train(data="data.yaml", epochs=100, batch=8, device="cuda", imgsz=640)