import torch
h = torch.randn(1, 2, 5, requires_grad=True); print(h)
val,idx = h.max(1, keepdim=True)
print(idx)
z = torch.nn.functional.gumbel_softmax(h, tau=2, hard=False, dim=1); print(z)
z = torch.nn.functional.gumbel_softmax(h, tau=2, hard=True, dim=1); print(z)

outputs are:

tensor([[[-0.0259, -0.9393, -0.2825,  0.6466, -1.0658],[-0.6078,  1.2127,  0.1509, -0.9749, -1.4952]]], requires_grad=True)
tensor([[[0, 1, 1, 0, 0]]])
tensor([[[0.7145, 0.1478, 0.8365, 0.5046, 0.4407],[0.2855, 0.8522, 0.1635, 0.4954, 0.5593]]], grad_fn=)
tensor([[[1., 0., 0., 1., 0.],[0., 1., 1., 0., 1.]]], grad_fn=)

the softmax procedure looks weird…