Input: - now_num: int - target_num: int Output: - indices: tensor[target_num]
(now_num, target_num, device='cuda')
| 695 | |
| 696 | # repeat them |
| 697 | def get_indices_for_repeat(now_num, target_num, device='cuda'): |
| 698 | """ |
| 699 | Input: |
| 700 | - now_num: int |
| 701 | - target_num: int |
| 702 | Output: |
| 703 | - indices: tensor[target_num] |
| 704 | """ |
| 705 | out_indice = [] |
| 706 | base_indice = torch.arange(now_num).to(device) |
| 707 | multiplier = target_num // now_num |
| 708 | out_indice.append(base_indice.repeat(multiplier)) |
| 709 | residue = target_num % now_num |
| 710 | out_indice.append(base_indice[torch.randint(0, |
| 711 | now_num, (residue, ), |
| 712 | device=device)]) |
| 713 | return torch.cat(out_indice) |
| 714 | |
| 715 | if self.dn_batch_gt_fuse: |
| 716 | raise NotImplementedError |