File size: 533 Bytes
78d2329 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | import torch
def batched_select(data, indices):
"""
Select data[i, indices[i]] for each batch element i.
Args:
data: [B, N, ...] input tensor
indices: [B, K] indices for each batch element
"""
assert data.shape[0] == indices.shape[0], f"Batch size mismatch {data.shape[0]} vs {indices.shape[0]}"
assert indices.dim() == 2, f"indices should be 2D, got {indices.shape}"
B = data.shape[0]
batch_idx = torch.arange(B, device=data.device)[:, None]
return data[batch_idx, indices]
|