src/transformers/models/jukebox/modeling_jukebox.py/JukeboxAttention/_slice_cache
class JukeboxAttention:    def _slice_cache(self, start, end=None):
        self.cache["key"] = self.cache["key"][:, start:end]
        self.cache["value"] = self.cache["value"][:, start:end]