`convert_ids_to_logits()` does not work with single-channel IDs

[ATATC]

When the input is single-channel, ids.unsqueeze(1).long() adds a redundant dimension.