Solved by albanD in post #6 
              
              
                There you go, this will be much faster  
Note that I use .narrow(-1, 0, x_size-1) out of habit but [:, :-1] works as well if you prefer that notation. 
import torch

A = torch.zeros((3, 3, 3), dtype = torch.float)
X = torch.tensor([[0, 1, 2, 0, 1, 0], [1, 0, 0, 2, 1, 1], [0, 0, 2, 2, 1, 1]])
for a, …