How to Train YOLOv7 on a Custom Dataset

I am following the blog Roboflow’s ‘How to Train YOLOv7 on a Custom Dataset’ for an weed identification model (Object detection). When I ran my dataset on the google colab notebook provided, it gave the following error. Please help me with this.

Hi, that’s apparently a known issue in the current implementation of YOLOv7.

This Stack Overflow post had the same issue. They say you can fix it by replacing the line in utils/loss.py line 685 to:

from_which_layer.append((torch.ones(size=(len(b),)) * i).to('cuda'))

also add a line after 756:

fg_mask_inboxes = fg_mask_inboxes.to(torch.device('cuda'))

Awesome, thank you! @stellasphere - I’ll send you an email as well with a personal thank you.

I’ve submitted a PR to the YOLOv7 maintainers with the fix to line 685 and the line added after line 756. I found after making the suggested changes from @stellasphere / Stack Overflow, the training runs fine.