learn.property_prediction.init_loss_fn

Contents

learn.property_prediction.init_loss_fn#

init_loss_fn(error_fn)[source]#

Returns a loss function to optimize model parameters.

Signature of error_fn:

error = error_fn(predictions, batch, mask)

where mask has the same shape as species to mask padded particles.

Parameters:

  • model – Molecular property prediction model (Haiku apply_fn).

  • error_fn – Error model quantifying the discrepancy between predictions and respective targets.

Contents