Some of the code may be wrong.

[Mike001-wq]

Should lines 234 and 235 in utils/general_utils.py be swapped? First record the val with the lowest probability, then replace it with the value with the highest probability. `

        weight[weight == val] = unique_vals[counts.argmax()]

        self.tail_locs[val.item()] = torch.where(weight == val)[0].detach().cpu()`