The issue is raised by the shape mismatch in:
    print(prob.shape, "---", label.shape)
    TP = (prob & label).sum(1).float()

as also seen in your print statements:
torch.Size([128, 258]) --- torch.Size([128])

so you would have to make sure the tensors have the same shape or can be broadcasted, if needed.