# How to Train YOLOv7 on a Custom Dataset

**URL:** <https://discuss.roboflow.com/t/how-to-train-yolov7-on-a-custom-dataset/1172>\
**Category:** 🤝  Community Help\
**Created:** [December 10, 2022, 6:29am UTC](https://discuss.roboflow.com/t/how-to-train-yolov7-on-a-custom-dataset/1172 "2022-12-10T06:29:40Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Akhileshsharma](https://yyz1.discourse-cdn.com/flex029/user_avatar/discuss.roboflow.com/akhileshsharma/32/510_2.png) [@Akhileshsharma](https://discuss.roboflow.com/u/Akhileshsharma)\
**Post date:** [December 10, 2022, 6:29am UTC](https://discuss.roboflow.com/t/how-to-train-yolov7-on-a-custom-dataset/1172/1 "2022-12-10T06:29:40Z")

</div>

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.

 ![error](https://canada1.discourse-cdn.com/flex029/uploads/roboflow1/original/1X/bac8a3bdf00a2306dc33e9c570a6b94a79130584.jpeg)

---

<div class="post-metadata">

**Author:** ![leo](https://yyz1.discourse-cdn.com/flex029/user_avatar/discuss.roboflow.com/leo/32/2228_2.png) [@leo](https://discuss.roboflow.com/u/leo)\
**Post date:** [December 13, 2022, 8:56pm UTC](https://discuss.roboflow.com/t/how-to-train-yolov7-on-a-custom-dataset/1172/2 "2022-12-13T20:56:33Z")

</div>

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

This [Stack Overflow](https://stackoverflow.com/a/74383094) post had the same issue. They say you can fix it by replacing the line in utils/loss.py line 685 to:

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

```

also add a line after 756:

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

```

---

<div class="post-metadata">

**Author:** ![Mohamed](https://yyz1.discourse-cdn.com/flex029/user_avatar/discuss.roboflow.com/mohamed/32/62_2.png) [@Mohamed](https://discuss.roboflow.com/u/Mohamed)\
**Post date:** [December 21, 2022, 5:23am UTC](https://discuss.roboflow.com/t/how-to-train-yolov7-on-a-custom-dataset/1172/3 "2022-12-21T05:23:22Z")

</div>

> [@leo](#):
>
> `fg_mask_inboxes = fg_mask_inboxes.to(torch.device('cuda'))`

Awesome, thank you! @leo - 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 @leo / Stack Overflow, the training runs fine.

- [Fix issues in loss.py by mo-traor3-ai · Pull Request #1283 · WongKinYiu/yolov7 · GitHub](https://github.com/WongKinYiu/yolov7/pull/1283)

 ![image](https://canada1.discourse-cdn.com/flex029/uploads/roboflow1/original/1X/24c80131564b2f368b5ba714dcecd1f3cca486d2.png)