Solved by Zhaoyi-Yan in post #4 
              
              
                import torch
import torch.nn.functional as F

ann_num = 3
h, w = 8, 8
c = 2
a = torch.randn(h*w, c)

S_mat = torch.cdist(a, a)

i_ind = torch.tensor([[12, 14, 51], [11, 6, 32], [1, 5, 17]]).long() # 3 ann_num
j_ind = torch.tensor([[21, 60, 33], [13, 16, 12], [7, 55, 29]]).long()
print('S_mat')
print…