How to train Yolo model with GPU of 'NVIDIA GeForce RTX 3050 Laptop GPU

I'm newbie to object detection model, I'm trying to train yolo model in my device('NVIDIA GeForce RTX 3050 Laptop GPU), I am getting errors, I have image size of (1365, 1024, 3) , I tried with image size(1365, 1024, 3) then my system struck. then I change with size into (1020, 640) , after the error like Yolo training recommended with 2 GPUs was it my system for training yolo. please help me with that. this is my code.

Load a model

from ultralytics import YOLO

model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)
model.train(data='datasets/m_train.yaml', 
                epochs=8, 
                imgsz=(1020, 640), 
                batch = 12,
                optimizer = 'Adam',
                device=[0,1]
                )

Solution

ultralytics already uses Dataloader() behind the scenes so there is no need for a custom implementation (I don't even know if you really need it, for other custom implementation/logic you would implement specific callbacks for customization).

The truth is that you have quite big image sizes in conjunction with big batch sizes, 3050Ti has only 4 GB of VRAM memory, so you are very likely to easily run into OOM issues.

You could try to start with (600x600) and a batch_size of 2 and then gradually increase the batch_size to (3,4 etc.)