Thanks,
I know that I can use scatter to create a one hot vector or something like that and in fact I have already done that:
idx = idx.data.new(input.size()).zero_().scatter_(-1, idx.view(-1, 1), 1.0)

but that problem is the output of the argmax will be fed into an embedding layer which doesnt accept one hot vectors, so in all cases I will need a single scalar value as the output.
I could have used the gumbel softmax instead, but it also has the same issue, it outputs a hot vector.