def sample_vectors(samples: torch.Tensor, num: int) -> torch.Tensor:
    """Sample vectors according to the given number."""
    num_samples, device = samples.shape[0], samples.device

    if num_samples >= num:
        indices = torch.randperm(num_samples, device=device)[:num]
    else:
        indices = torch.randint(0, num_samples, (num, ), device=device)

    return samples[indices]
Sample vectors according to the given number.
mmpretrain/models/utils/vector_quantizer.py