Solved by ptrblck in post #2 
              
              
                I’m sure there are better approaches and you could speed it up more, but you could use index_put_ with the accumulate=True argument instead of the for loop and precompute some indices. 
Here is a code snippet: 
def fun1():
    result = torch.full(( 32, 256, 256), -1.)
    data = torch.full(( 32, 400…