def load(
        self,
        booster: Booster,
        load_dir: str,
        model: nn.Module = None,
        ema: nn.Module = None,
        optimizer: Optimizer = None,
        lr_scheduler: _LRScheduler = None,
        sampler=None,
        strict: bool = False,
        include_master_weights: bool = False,
    ) -> tuple[int, int]:
        """
        Load a checkpoint.

        Args:
            booster (Booster): The Booster object.
            load_dir (str): The directory to load the checkpoint from.
            model (nn.Module): The model to load the checkpoint into.
            ema (nn.Module): The EMA model to load the checkpoint into.
            optimizer (Optimizer): The optimizer to load the checkpoint into.
            lr_scheduler (_LRScheduler): The learning rate scheduler to load the checkpoint into.
            sampler: The sampler to load the checkpoint into.

        Returns:
            tuple[int, int]: The epoch and step of the checkpoint.
        """
        assert os.path.exists(load_dir), f"Checkpoint directory {load_dir} does not exist"
        assert os.path.exists(os.path.join(load_dir, "running_states.json")), "running_states.json does not exist"

        running_states = load_json(os.path.join(load_dir, "running_states.json"))
        if model is not None:
            booster.load_model(
                model,
                _search_valid_path(os.path.join(load_dir, "model")),
                strict=strict,
                low_cpu_mem_mode=False,
                num_threads=32,
            )
        if ema is not None:
            if os.path.exists(os.path.join(load_dir, "ema.safetensors")):
                ema_state_dict = load_file(os.path.join(load_dir, "ema.safetensors"))
            else:
                ema_state_dict = torch.load(os.path.join(load_dir, "ema.pt"), map_location=torch.device("cpu"))
            # ema is not boosted, so we don't use booster.load_model
            ema.load_state_dict(ema_state_dict, strict=strict, assign=True)

        if optimizer is not None:
            booster.load_optimizer(
                optimizer, os.path.join(load_dir, "optimizer"), low_cpu_mem_mode=False, num_threads=32
            )
            if include_master_weights:
                master_state_dict = load_file(os.path.join(load_dir, "master.safetensors"))
                load_master_weights(model, optimizer, master_state_dict)
        if lr_scheduler is not None:
            booster.load_lr_scheduler(lr_scheduler, os.path.join(load_dir, "lr_scheduler"))
        if sampler is not None:
            sampler.load_state_dict(torch.load(os.path.join(load_dir, "sampler")))

        dist.barrier()

        return (running_states["epoch"], running_states["step"])

        Load a checkpoint.

        Args:
            booster (Booster): The Booster object.
            load_dir (str): The directory to load the checkpoint from.
            model (nn.Module): The model to load the checkpoint into.
            ema (nn.Module): The EMA model to load the checkpoint into.
            optimizer (Optimizer): The optimizer to load the checkpoint into.
            lr_scheduler (_LRScheduler): The learning rate scheduler to load the checkpoint into.
            sampler: The sampler to load the checkpoint into.

        Returns:
            tuple[int, int]: The epoch and step of the checkpoint.
        
opensora/utils/ckpt.py