My question is the way of row-wise masking in batch manner.
I wrote non-batch case in the following.
import torch

n_dim0, n_dim1 = 4, 100
threshold = 0.5

storage = torch.rand((n_dim0, n_dim1))
mask = (threshold < storage)

print('storage:', storage.shape)  # [4, 100]
print('mask:   ', mask.shape)     # [4, 100]

mask_lens = torch.sum(mask, axis=1)   # [4]
max_mask_lens = torch.max(mask_lens)  # [1]

masked = torch.zeros((n_dim0, max_mask_lens))
masked[0, 0:mask_lens[0]] = storage[0, mask[0]]  #
masked[1, 0:mask_lens[1]] = storage[1, mask[1]]  # How to make these lines
masked[2, 0:mask_lens[2]] = storage[2, mask[2]]  # batch operation?
masked[3, 0:mask_lens[3]] = storage[3, mask[3]]  #

print('masked: ', masked.shape)  # [4, X]  (X is around 50)

I wish someone has great idea.
Thank you.