[Solved] pytorch Error: KeyError: tensor(2, device=‘cuda:0‘)

Solution: convert the index strength of the array to int

Example:

# original
length = [id2tags[i.item()] for i in y[0].cpu() if i.item() > 0]
# Modify to:
length = [id2tags[int(i.item())] for i in y[0].cpu() if i.item() > 0]

Read More: