INNER CODE UNIT · Python
calc_aux_loss
lucidrains/vector-quantize-pytorch · vector_quantize_pytorch/binary_mapper.py:75
def calc_aux_loss(
self,
logits,
reduce_aux_kl_loss = True
):
logits, inverse_pack_lead_dims = pack_with_inverse(logits, '* bits')
kl_div = self.bits * NAT - self.binary_entropy(logits)
aux_kl_loss = F.relu(kl_div - self.kl_loss_threshold)
if reduce_aux_kl_loss:
return aux_kl_loss.mean()
return inverse_pack_lead_dims(aux_kl_loss, '*')
def log_prob(
self,
logits,
*,