cube/cube3d/model/transformers/cache.py

from dataclasses import dataclass

import torch


@dataclass
class Cache:
    key_states: torch.Tensor
    value_states: torch.Tensor