Solved by ptrblck in post #2 
              
              
                It looks like your index is missing a row or your array has one additional row (row2 based on the result). 
A combination of scatter_add_ and gather would work: 
array = torch.tensor([[0., 1, 2, 3, 4, 5, 6],
                      [7, 8, 9, 10, 11, 12, 13],
                      [21, 22, 23, 24, 25, …