albanD:

logits = self.mask * logits + (1 - self.mask) * CategoricalMasked.LOGIT_INF_VALUE



that shouldn’t be faster than: torch.lerp(LOGIT_INF_TENSOR, logits, self.mask)



 AdilZouitine:

p_log_p = torch.where(self.mask, p_log_p, torch.tensor(0., device=self.logits.device)



p_log_p *= self.mask