@Varal7  Thanks for your reply. it is very helpful. I have studied the gumbel_softmax and change the code in the model forward function as the following:
distance = torch.norm(distance,dim=1)
out = 1 - torch.nn.functional.gumbel_softmax(dist)
Because I am going to find minimum, I subtract the output of gumbel_softmax by 1.
I am not sure my implementation is correct. Could you give me more hint on the issueto figure out problem of my work.