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