Solved by andreaskoepf in post #2 
              
              
                You could run a loop over the 100 entries in the first dimension, e.g.: 
import torch

#x = torch.rand(100, 64, 256, 256)
x = torch.rand(100, 64, 3, 3) # reduced for testing

def max_exclusive(x, i):
    a,b = x[:i], x[i+1:]
    if a.nelement() == 0:
        return torch.max(b, 0)[0]
    if b.neleme…