Hi @smth, Will backprop work fine through a tensor output from torch.gather? Reason I became suspicious is because I found out that torch.gather on an expanded tensor (i.e., output of torch.expand) resulted in a larger storage area than the original tensor (due to the expansion I am guessing).

Tensor Y: Shape = (12, 15, 512), Storage = 368704 bytes
Tensor indices: Shape = (12, 15, 5, 512), Storage = 664 bytes
Y_expanded = Y.unsqueeze(2).expand(-1, -1, 15, -1)
Tensor Y_expanded: Shape = (12, 15, 15, 512), Storage = 368704 bytes

Up until this point, size of Y_expanded == size of Y

Y_selected = torch.gather(Y_expanded, dim=2, index=m_indices)
Tensor Y_selected: Shape = (12, 15, 5, 512), Storage = 1843264 bytes

However, after I select from Y_expanded, the storage size increases, implying a copy into new storage. Is there a way I can avoid the increase in size? Also, will backprop flow back through Y if I use this method?
Thanks.