Solved by ptrblck in post #2 
              
              
                This should work: 
a = torch.tensor([[[[160, 56, 20],
                    [54, 6, 97],
                    [65, 119, 56]],
                 
                   [[185, 13, 90],
                    [10, 6, 220],
                    [67, 88, 4]],

                   [[210, 92, 74],
                    …