pn_loss

probly.losses.torch.pn_loss(model: Module, x_in: Tensor, y_in: Tensor, x_ood: Tensor) → Tensor[source]

Dirichlet Prior Network loss based on [MG18].

Combines KL divergence to sharp in-distribution targets and flat out-of-distribution targets, with an additional cross-entropy term for classification stability.

Parameters:

  • model – Network mapping inputs to Dirichlet concentration parameters.

  • x_in – In-distribution inputs, shape (B, …).

  • y_in – In-distribution class labels, shape (B,).

  • x_ood – Out-of-distribution inputs, shape (B_ood, …).

Returns:

Scalar paired ID+OOD Prior Networks loss.