| 29 | |
| 30 | |
| 31 | class ProjectReadout(nn.Module): |
| 32 | def __init__(self, in_features, start_index=1): |
| 33 | super(ProjectReadout, self).__init__() |
| 34 | self.start_index = start_index |
| 35 | |
| 36 | self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), nn.GELU()) |
| 37 | |
| 38 | def forward(self, x): |
| 39 | readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index :]) |
| 40 | features = torch.cat((x[:, self.start_index :], readout), -1) |
| 41 | |
| 42 | return self.project(features) |
| 43 | |
| 44 | |
| 45 | class Transpose(nn.Module): |