Thank you for you reply @ptrblck !
I have seen theirs but it requires more modifications to my code… Can I use it as my model but train it with my own train.py code?
I have also tried to use both their model and their loss function existed in train.py (class ComputeLoss in loss.py) but I haven’t found documentation on what the input ‘targets’ to this function should be.
What I mean to do is:
for epoch in tqdm(range(train_params.num_epochs)):
    model.train()
    for i, (images, bboxs, labels) in enumerate(train_loader):
        preds = YOLO_model(images)
        loss = compute_loss(preds, targets)
        loss.backward()


What should I give in ‘targets’?
And how do I separate ‘preds’  to labels and bboxes?

Thank you!