Create mask tensor from lengths without for loop

1

Hi everyone,

I would like to create a mask tensor filled with ones where values should be masked, that masks a part of the input sequences depending on their lengths, used with a Transformer architecture. lengths is a 1D tensor containing the length of all input sequences in a batch. The code works but I would like to do it without a for loop. Any ideas?

Thanks a lot

torch::Tensor create_mask(torch::Tensor &lengths,
                                            const int percentage)
{
  int max_len = torch::max(lengths).item<int>();
  auto mask_src = torch::zeros({lengths.size(0), max_len});

  for(int i = 0; i < lengths.size(0); i++)
  {
    int current_length = lengths.index({i}).item<int>();

    int nb_masked = current_length * percentage / 100;

    int start_index =
        torch::randint(0, current_length - nb_masked, {1}).item<int>();

    mask_src.index_put_({i, Slice(start_index, start_index + nb_masked)}, 1);
  }

  return mask_src;
}