Dimension inconsistency of the token labels

[lingeringlight]

Hi, I am curious about the problem of dimension inconsistency.
(1) The shape of "score_map" that generated in generate_label.py is [2, 5, H, W], but the dimension of score_maps seems to be [2, H, W, 5] in "score_maps[-1,0,0,5]=target " (line 97 of TokenLabeling/tlt/data/dataset.py )
(2) The dimension of "label_maps_topk" in line 54 of TokenLabeling/tlt/data/mixup.py is [batch_size, 3, H, H, 5], but I cannot find how to transform "score_maps" to "label_maps_topk", and what information is stored in the 0, 1, 2 dimension of "label_maps_topk", respectively.
This problem has also been mentioned in the Issue #9