(feat, ind, mask=None)
| 10 | return y |
| 11 | |
| 12 | def _gather_feat(feat, ind, mask=None): |
| 13 | dim = feat.size(2) |
| 14 | ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim) |
| 15 | feat = feat.gather(1, ind) |
| 16 | if mask is not None: |
| 17 | mask = mask.unsqueeze(2).expand_as(feat) |
| 18 | feat = feat[mask] |
| 19 | feat = feat.view(-1, dim) |
| 20 | return feat |
| 21 | |
| 22 | def _transpose_and_gather_feat(feat, ind): |
| 23 | feat = feat.permute(0, 2, 3, 1).contiguous() |
no test coverage detected