You could probably avoid the last loop using:
idx = activation.sum(1) == counter
features_list = torch.cat((features_list, outputs[idx]), 0)

However, as I don’t have dummy inputs and the expected results, I couldn’t really test it, so you would need to verify it.
I don’t think you could avoid the dist_tensor calculation, as it would be recalculated for the current output sample, if I understand the use case correctly.