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