| 181 | |
| 182 | |
| 183 | class QueryNetwork(nn.Module): |
| 184 | def __init__(self): |
| 185 | super(QueryNetwork, self).__init__() |
| 186 | # TODO: Add proper non-linearity. |
| 187 | self.model = nn.Sequential( |
| 188 | nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 288), nn.ReLU(), nn.Linear(288, 288) |
| 189 | ) |
| 190 | |
| 191 | def forward(self, query): |
| 192 | out = self.model(query) |
| 193 | return out.reshape(-1, 4, 72) |
| 194 | |
| 195 | |
| 196 | class SpatialBasis: |