Solved by albanD in post #6 
              
              
                Sure something like that will work: 
import torch

inp = torch.tensor([[[[0., 0., 1.],
          [0., 0., 1.],
          [0., 1., 0.],
          [0., 1., 0.],
          [0., 0., 1.],
          [0., 0., 1.],
          [0., 1., 0.],
          [0., 0., 1.],
          [1., 0., 0.],
          [0., 0., 1.…