| 465 | |
| 466 | |
| 467 | class T5Block(torch.nn.Module): |
| 468 | def __init__(self, model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias, dtype, device): |
| 469 | super().__init__() |
| 470 | self.layer = torch.nn.ModuleList() |
| 471 | self.layer.append(T5LayerSelfAttention(model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias, dtype, device)) |
| 472 | self.layer.append(T5LayerFF(model_dim, ff_dim, dtype, device)) |
| 473 | |
| 474 | def forward(self, x, past_bias=None): |
| 475 | x, past_bias = self.layer[0](x, past_bias) |
| 476 | x = self.layer[-1](x) |
| 477 | return x, past_bias |
| 478 | |
| 479 | |
| 480 | class T5Stack(torch.nn.Module): |