Input: - now_num: int - target_num: int Output: - indices: tensor[target_num]
(now_num, target_num, device='cuda')
| 2734 | |
| 2735 | # repeat them |
| 2736 | def get_indices_for_repeat(now_num, target_num, device='cuda'): |
| 2737 | """ |
| 2738 | Input: |
| 2739 | - now_num: int |
| 2740 | - target_num: int |
| 2741 | Output: |
| 2742 | - indices: tensor[target_num] |
| 2743 | """ |
| 2744 | out_indice = [] |
| 2745 | base_indice = torch.arange(now_num).to(device) |
| 2746 | multiplier = target_num // now_num |
| 2747 | out_indice.append(base_indice.repeat(multiplier)) |
| 2748 | residue = target_num % now_num |
| 2749 | out_indice.append(base_indice[torch.randint(0, |
| 2750 | now_num, (residue, ), |
| 2751 | device=device)]) |
| 2752 | return torch.cat(out_indice) |
| 2753 | |
| 2754 | if self.dn_batch_gt_fuse: |
| 2755 | raise NotImplementedError |