The whole function can be simplified to this:
def AdaReLU(x):                    
    max_val, _ = torch.max(x, 3)   
    mask = x < max_val.expand_as(x)
    mask.detach_()                 
    x = x.clone()                  
    x[mask] = 0                    
    return F.relu(x)               

The detach_() call shouldn’t be necessary, but I’ve found a bug there