Solved by LeviViana in post #2 
              
              
                This way: 
import torch

population = torch.tensor([
                   [6, 8, 10,  7, 5, 0, 0],
                   [5, 7,  9, 11, 8, 7, 3],
                   [4, 6,  6,  8, 5, 4, 1]
             ]
)

moves = torch.tensor([
                   [0, 0,  0,  0, 2, 0, 0],
                   [0, 0, -1,  …