820f6f8
1
2
3
4
5
6
7
8
9
10
from dataclasses import dataclass from typing import * import torch @dataclass class Transformer1DModelOutput: sample: torch.FloatTensor aux_loss: torch.FloatTensor