| 358 | |
| 359 | |
| 360 | class NonEquivariantMaxPool(nn.Module): |
| 361 | def __init__(self, dim=-1): |
| 362 | super(NonEquivariantMaxPool, self).__init__() |
| 363 | self.dim = dim |
| 364 | |
| 365 | def forward(self, x): |
| 366 | ''' |
| 367 | x: point features of shape [B, C, N] |
| 368 | ''' |
| 369 | return torch.max(x, dim=self.dim, keepdim=True)[0] |
| 370 | |
| 371 | |
| 372 | class NonEquivariantStdFeature(nn.Module): |