Add more detailed docstring to knn predict

Open IgorSusmelj opened this issue 2 years ago • 2 comments

We have here the knn predict code: https://github.com/lightly-ai/lightly/blob/df004d288fa07816e55975fa0af92868f9867436/lightly/utils/benchmarking/knn.py#L8

The problem is that the docstring doesn't tell us much about the returned values. One would expect a list of labels or a tensor with probabilities for all the classes. But we return a NxD shape tensor that contains the class indices sorted by probability. We should clarify that in the docstring.