import torch
torch.cuda.empty_cache()
from ultralytics import YOLO

model = YOLO("yolov8n.pt")  
model.train(data="/home/capstone_002_02/Lina/data/data.yaml",  
            epochs=50,                   
            batch=2,               
            device='cuda')              
model.save("/home/capstone_002_02/Lina/parameter/5_9_yolo_model_2.pth")
results = model.evaluate(data="/home/capstone_002_02/Lina/data/data.yaml", imgsz=640, batch=2)
print('Model performance:', results)