File size: 208 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 | from __future__ import annotations
from dataclasses import dataclass
import torch
@dataclass(frozen=True)
class TopKOutput:
weights: torch.Tensor
ids: torch.Tensor
router_logits: torch.Tensor
|