masked_for_selection_token doesn't match matching_logits_ #12
Description
At line 690 in model.py, masked_for_selection_token corresponds to the range [1:] which ignores the start token while latent corresponds to the range [0:end) which take the start into account. Their positions doesn't match
