    def _load_model(self, checkpoint_dir, from_multi_gpu=False):
        """Load model from checkpoint. If a folder is given, it will
        load the latest checkpoint in checkpoint_dir. If a path is given
        it will load the checkpoint specified by checkpoint_path.
        **Only use this method after** ``accelerator.prepare()``.
        """
        if os.path.isdir(checkpoint_dir):
            if "epoch" in checkpoint_dir and "step" in checkpoint_dir:
                checkpoint_path = checkpoint_dir
            else:
                # Load the latest accelerator state dicts
                ls = [
                    str(i)
                    for i in Path(checkpoint_dir).glob("*")
                    if not "audio" in str(i)
                ]
                ls.sort(
                    key=lambda x: int(x.split("/")[-1].split("_")[0].split("-")[-1]),
                    reverse=True,
                )
                checkpoint_path = ls[0]
            accelerate.load_checkpoint_and_dispatch(
                self.accelerator.unwrap_model(self.model),
                os.path.join(checkpoint_path, "pytorch_model.bin"),
            )
            return str(checkpoint_path)
        else:
            # Load old .pt checkpoints
            if self.cfg.model.generator in [
                "bigvgan",
                "hifigan",
                "melgan",
                "nsfhifigan",
            ]:
                ckpt = torch.load(
                    checkpoint_dir,
                    map_location=(
                        torch.device("cuda")
                        if torch.cuda.is_available()
                        else torch.device("cpu")
                    ),
                )
                if from_multi_gpu:
                    pretrained_generator_dict = ckpt["generator_state_dict"]
                    generator_dict = self.model.state_dict()

                    new_generator_dict = {
                        k.split("module.")[-1]: v
                        for k, v in pretrained_generator_dict.items()
                        if (
                            k.split("module.")[-1] in generator_dict
                            and v.shape == generator_dict[k.split("module.")[-1]].shape
                        )
                    }

                    generator_dict.update(new_generator_dict)

                    self.model.load_state_dict(generator_dict)
                else:
                    self.model.load_state_dict(ckpt["generator_state_dict"])
            else:
                self.model.load_state_dict(torch.load(checkpoint_dir)["state_dict"])
            return str(checkpoint_dir)
Load model from checkpoint. If a folder is given, it will
load the latest checkpoint in checkpoint_dir. If a path is given
it will load the checkpoint specified by checkpoint_path.
**Only use this method after** ``accelerator.prepare()``.
models/codec/codec_inference.py