| 16 | |
| 17 | |
| 18 | class AddReadout(nn.Module): |
| 19 | def __init__(self, start_index=1): |
| 20 | super(AddReadout, self).__init__() |
| 21 | self.start_index = start_index |
| 22 | |
| 23 | def forward(self, x): |
| 24 | if self.start_index == 2: |
| 25 | readout = (x[:, 0] + x[:, 1]) / 2 |
| 26 | else: |
| 27 | readout = x[:, 0] |
| 28 | return x[:, self.start_index :] + readout.unsqueeze(1) |
| 29 | |
| 30 | |
| 31 | class ProjectReadout(nn.Module): |