train.PredictionWriter

class train.PredictionWriter(*args: Any, **kwargs: Any)

Callback to store prediction results using torch.save.

Parameters:

• output_dir – output directory for prediction files

• write_interval – can be one of [“batch”, “epoch”, “batch_and_epoch”]

• write_idx – Write molecular ids for all atoms. This is needed for atomic properties like forces.