Method get_rope_cache

lit_gpt/model.py:193–201

(self, device: Optional[torch.device] = None)

        return cls(Config.from_name(name, **kwargs))

    def get_rope_cache(self, device: Optional[torch.device] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        return build_rope_cache(
            seq_len=self.max_seq_length,
            n_elem=self.config.rope_n_elem,
            device=device,
            condense_ratio=self.config.rope_condense_ratio,
            base=self.config.rope_base,
            dtype=torch.bfloat16
        )

    def build_rope_cache(self, idx: torch.Tensor) -> RoPECache:
        return build_rope_cache(