Solved by melike in post #8 
              
              
                I found a way to do it without a for loop, combining the approaches of @tom, @ptrblck and with some support from friends, 
import torch
from torch_scatter import scatter_max

class AdvMaxify(torch.autograd.Function):
    
    @staticmethod
    def forward(ctx, inp, CCs):
        indexes = torch.uniq…