Solved by KFrank in post #3 
              
              
                Hi lkp411! 

Yes, the idea would be to construct the combinations of a 1D vector 
of indices and then index into your multidimensional tensor, swapping 
the dimensions around appropriately. 
Like this: 
>>> import torch
>>> torch.__version__
'1.9.0'
>>> a = torch.tensor([[2, 5, 6], [7, 9, 4]])
>>> a…