I am trying to take the mean across a tensor, data, having shape Batch_Size * Seq_len * Embed_dim and using a list of specified lengths across which I need to take an average. So my output should be of shape Batch_size * input_lens.shape[1] * Embed_dim where input_lens is my tensor containing the lengths of shape Batch_size * num_lengths. So far what I have got is something like :
for i, batch in enumerate(input_lens) :
        	base_length = 0
        	for j,length in enumerate(batch) :
        		mean_vec[i,j,:] = torch.mean(data[i,base_length:length,:], dim=0)
        		base_length += length.data

Is there a vectorised Pytorch way of doing this?
Thanks,
Umang