File size: 175 Bytes
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