Confused by this LogitMask

[JJASMINE22]

rqvae/models/rqtransformer/primitives.py
class LogitMask(nn.Module):
def init(self, vocab_size: Iterable[int], value=-1e6):
super().init()

    self.vocab_size = vocab_size
    self.mask_cond = [vocab_size[0]]*len(vocab_size) != vocab_size
    self.value = value

def forward(self, logits: Tensor) -> Tensor:
    if not self.mask_cond:
        return logits
    else:
        for idx, vocab_size in enumerate(self.vocab_size):
            logits[:, idx, vocab_size:].fill_(-float('Inf'))
        return logits

The logits mentioned in LogitMask should probably be expressed as logits[:, idx, vocab_size:] = logits[:, idx, vocab_size:].fill_(-float('Inf'))