To elaborate more, fundamentally I am looking for some operation that lets you slice based of a LongTensor of start or end points. So like a “assign_slices” method on tensors which takes in (starts, ends, new_val, dim)
>>> a = torch.tensor([[1,2,3],[4,5,6]])
>>> a.assign_slices(torch.tensor([2,1], None, 0, dim=0)
torch.tensor([[1, 2, 0], [4, 0, 0])
>>> a.assign_slices(None, torch.tensor([2,1]), 99, dim=0)
torch.tensor([[99, 99, 3], [99, 5, 6])
>>> a.assign_slices( torch.tensor([1, 0]), torch.tensor([2, 2]), 0, dim=0)
torch.tensor([[1, 0, 3], [0, 0, 6])

Does something like this exist?
If it does then I can just make everything after the input_lens be 0 and then sum as normal to get my average.
I found myself suddenly again wanting this same functionality when trying to make a model predict  spans of tokens where the end of the span was conditioned to be after the start (so basically the starts was a long tensor, and everything before it I wanted to set a large negative value).