braintools.trainer.load_checkpoint

braintools.trainer.load_checkpoint(filepath, map_location=None)

Load a checkpoint from disk.

Parameters:

• filepath (str) – Path to the checkpoint file.

• map_location (str | None) – Device to map tensors to (not used in JAX, included for API compatibility).

Returns:

Loaded state dictionary.

Return type:

Dict[str, Any]

Examples

>>> state = load_checkpoint('checkpoint.ckpt')
>>> model.load_state_dict(state['model'])