abeardear / pytorch-YOLO-v1

an experiment for yolo-v1, including training and testing.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

IndexError: invalid index of a 0-dim tensor

zjykzj opened this issue · comments

Met this question when using predict.py. Error code location is located at

        i = order[0]
        keep.append(i)

after check the code, I found the key is to update squeeze() use

        # before
        ids = (ovr<=threshold).nonzero().squeeze()
        # after
        ids = (ovr<=threshold).nonzero().squeeze(-1)

this change can solve the problem, when ids shape == (1,1), then squeeze(-1) return (1)