tomodrgn.commands.train_nn.save_checkpoint

save_checkpoint(model: FTPositionalDecoder | DataParallelPassthrough, lat: Lattice, optim: Optimizer, epoch: int, norm: tuple[float, float], config_path: str, amp: bool, out_mrc: str, out_weights: str) → None

Evaluate the model with the currently loaded weights; save the resulting consensus volume and model.

Parameters:

• model – FTPositionalDecoder model to be evaluated

• lat – lattice object specifying coordinates to evaluate

• optim – optimizer object, to be saved into model pkl

• epoch – current epoch number, 0-indexed, saved into model pkl

• norm – mean and standard deviation of dataset by which input dataset was normalized, to be unscaled following volume generation

• config_path – path to trained model config.pkl

• amp – Enable or disable use of mixed-precision model inference

• out_mrc – name of the output mrc file to save consensus reconstruction

• out_weights – name of the output pkl file to save model state_dict, optimizer state_dict

Returns:

None