Quick fix to avoid error when not enough proposals post augmentation

[artemisp]

#181

[kusstox]

If missing is longer than topk_logits.shape[-1], you will run in the same issue again. To fix it, we should embed it into a while clause

[Parul-Gupta]

The following code snippet covers the case where missing is longer than topk_logits.shape[-1]

if topk > topk_logits.shape[-1]:
    missing =  topk-topk_logits.shape[-1]
    num_repeats = missing//topk_logits.shape[-1]
    if num_repeats>0:
        missing = missing - num_repeats*topk_logits.shape[-1]
        repeat_pad1 = topk_logits.clone().detach().repeat(1, num_repeats)
        repeat_pad = topk_logits.clone().detach()[:,-missing:]
        topk_logits = torch.cat([topk_logits, repeat_pad1, repeat_pad], dim=-1)

        repeat_pad1 = output_proposals.clone().detach().repeat(1, num_repeats, 1)
        repeat_pad = output_proposals.clone().detach()[:,-missing:, :]
        output_proposals = torch.cat([output_proposals, repeat_pad1, repeat_pad], dim=-2)

        repeat_pad1 = enc_outputs_coord_unselected.clone().detach().repeat(1, num_repeats, 1)
        repeat_pad = enc_outputs_coord_unselected.clone().detach()[:,-missing:, :]
        enc_outputs_coord_unselected = torch.cat([enc_outputs_coord_unselected, repeat_pad1, repeat_pad], dim=-2)

        repeat_pad1 = output_memory.clone().detach().repeat(1, num_repeats, 1)
        repeat_pad = output_memory.clone().detach()[:,-missing:, :]
        output_memory = torch.cat([output_memory, repeat_pad1, repeat_pad], dim=-2)
    else:
        repeat_pad = topk_logits.clone().detach()[:,-missing:]
        topk_logits = torch.cat([topk_logits, repeat_pad], dim=-1)

        repeat_pad = output_proposals.clone().detach()[:,-missing:, :]
        output_proposals = torch.cat([output_proposals, repeat_pad], dim=-2)

        repeat_pad = enc_outputs_coord_unselected.clone().detach()[:,-missing:, :]
        enc_outputs_coord_unselected = torch.cat([enc_outputs_coord_unselected, repeat_pad], dim=-2)

        repeat_pad = output_memory.clone().detach()[:,-missing:, :]
        output_memory = torch.cat([output_memory, repeat_pad], dim=-2)