Solved by ptrblck in post #2 
              
              
                You could create a bool mask and index the input tensor as: 
x = torch.tensor([[[ 1,  2,  3,  4],
                   [ 5,  6,  7,  8],
                   [ 9, 10, 11, 12]],

                  [[13, 14, 15, 16],
                   [17, 18, 19, 20],
                   [21, 22, 23, 24]]])

mask = torch…