nexuml_library.layers.loss.classification_loss

nexuml_library.layers.loss.classification_loss

Classification loss layer.

ClassificationLoss

Bases: LayerDefinition

Computes classification loss from logits and labels.

keys_in: [logits_key] keys_out: [loss_key] Reads labels from y TensorDict using label_key.

LogitsToClass

Bases: LayerDefinition

Converts class logits to predicted class integers via argmax.