Args:
checkpoint_file (`str` or `os.PathLike`):
Path to the checkpoint file to load. Can be either a safetensors or pickle (`.bin`) checkpoint.
map_location (`str` or `torch.device`, *optional*):
A `torch.device` object, string or a dict specifying how to remap storage locations. It
indicates the location where all tensors should be loaded.
weights_only (`bool`, *optional*, defaults to `False`):
If True, only loads the model weights without optimizer states and other metadata.