Thanks for replying.
I find that my statement is misleading. Sorry about that, let me make the problem more clear.
Suppose I have 300 vector store in a tensor called feature, which is in the shape [300, 1024]. And we also have an index tensor, which is in [300] size and indicates which vector should send into max() operation to get maximum tensor. so, index is something like [1,0,1,0,0,1,…,0].
The easiest way to do this is just one line code:
out = torch.max(feature[torch.where(index == 1)], 0)
But, we have a limit memory, and I guess this part feature[torch.where(index == 1)] will allocate a new memory, which is not acceptable.
So, I think there is another way: first, allocate space in memory for storage maximum result. Then, use a for-loop to compare each feature with result one-by-one.
out = torch.zeros([1024])
for i in range(300):
    if index[i] == 1:
        out = torch.max(out, feature[i])

So, ideally, we only need to allocate one [1024] size memory for out tensor, and this is acceptable by our gpu. (Both index and feature are tensors here, but we don’t want to create a big intermediate variable so we use for-loop)
However, I find that during the for-loop, each iteration, memory is still growing. I don’t know why, and how to avoid it.
We need  out tensor contains gradient for back-propagation. So I guess this is the result pytorch doesn’t release the memory after torch.max(out, feature[i]). But, ideally, we just need one same size tensor with out to remember where the max value come from, then we can do backpropagation. Is there a way to do that?