nvalchemi.training.default_training_fn#

nvalchemi.training.default_training_fn(model, batch)[source]#

Run a forward pass and prefix output keys with predicted_.

Parameters:

model (BaseModelMixin) – A wrapped MLIP whose __call__ returns model outputs.
batch (Batch) – Input batch of atomic graphs.

Returns:

Predictions keyed by predicted_<output_name> with None outputs omitted.

Return type:

dict[str, torch.Tensor]