thanks a lot, it just had one minor proble, we should use dsum = torch.abs(d).sum(-1) instead of d.sum(-1), because the sum of numbers might lead to zero, although they are not all zeros.
see the following e.g.
    T = torch.round(torch.rand(1,3,4)*10)
    4 3 6 3
    1 0 5 8
    1 10 4 8
    K = torch.cat((torch.round(torch.rand(1,3,4)*10),T, torch.zeros(1,3,4)),1)
    0 9 1 9
    3 3 4 7
    9 3 1 3
    4 3 6 3
    1 0 5 8
    1 10 4 8
    0 0 0 0
    0 0 0 0
    0 0 0 0
    D = torch.round(torch.rand(1,9,4)*10)
    3 6 7 2
    4 0 5 9
    4 2 5 10
    4 4 9 2
    2 0 2 6
    2 1 4 0
    1 4 8 3
    4 3 8 0
    2 3 9 10
    d = T.unsqueeze(2) - K.unsqueeze(1)
    (0 ,0 ,.,.) =
    4 -6 5 -6
    1 0 2 -4
    -5 0 5 0
    0 0 0 0
    3 3 1 -5
    3 -7 2 -5
    4 3 6 3
    4 3 6 3
    4 3 6 3

    (0 ,1 ,.,.) =
    1 -9 4 -1
    -2 -3 1 1
    -8 -3 4 5
    -3 -3 -1 5
    0 0 0 0
    0 -10 1 0
    1 0 5 8
    1 0 5 8
    1 0 5 8

    (0 ,2 ,.,.) =
    1 1 3 -1
    -2 7 0 1
    -8 7 3 5
    -3 7 -2 5
    0 10 -1 0
    0 0 0 0
    1 10 4 8
    1 10 4 8
    1 10 4 8

    dsum = d.sum(-1)
    (0 ,.,.) =
    -3 -1 0 0 2 -7 16 16 16
    -5 -3 -2 -2 0 -9 14 14 14
    4 6 7 7 9 0 23 23 23
    [torch.FloatTensor of size 1x3x9]

    loc = (dsum==0).nonzero()
    0 0 2
    0 0 3
    0 1 4
    0 2 5
    loc[:,-1]
    2
    3
    4
    5
    D[:,loc[:,-1],:]
    (0 ,.,.) =
    4 2 5 10
    4 4 9 2
    2 0 2 6
    2 1 4 0
    [torch.FloatTensor of size 1x4x4]