segmentation with yolov5

[rhsaeedy]

I'm following this tutorial to design a custom instance segmentation on OAK-D. I trained the custom model with only one class and used depthai online tool to convert it to the blob. When I use main.py from host-decoding folder I get the following error:

Traceback (most recent call last):
  File "~/depthai-experiments/gen2-yolo/host-decoding/main.py", line 170, in <module>
    boxes = non_max_suppression(output, conf_thres=conf_thresh, iou_thres=iou_thresh)
  File "~/depthai-experiments/gen2-yolo/host-decoding/util/functions.py", line 29, in non_max_suppression
    xc = prediction[..., 4] > conf_thres  # candidates
IndexError: index 0 is out of bounds for dimension 1 with size 0

I also tried other values instead of 4 in xc = prediction[..., 4] and got the same error.
I trained the NN with 640x640 image size. For now I really want to pass this error but my final goal is to be able to get the polygon masks that I used to label the data with for instance segmentation. I wonder if it is possible or not?
By the way, when I print the prediction.shape I get torch.Size([1, 10647, 0])