This is your setup doing it. I don’t think there is much that PyTorch can do for you here, because your operation is not inplace.
When you do pars*sample that requires grad. In turn the cat produces a tensor that requires grad and refers to pars*sample and the previous out as inputs. Then leaveTopK indexes into that and has out as input. etc.
One way to get around that is to keep the indexes into pars and operate with those and the samples (which do not require grad). Then in each iteration rebuild out = cat(pars[idxes-of-samples-you-kept] * samples-you-kept, pars * new_samples) It’s a bit more elaborate but allows PyTorch autograd to forget about the previous step (which it can only if your restart with pars rather than something computed from pars).
Best regards
Thomas
P.S.: You didn’t ask, but I think it’s canonical to just call the torch module torch.