Hey Thomas,
Thank you very much for youhelp. After going through the whole thing over and over again I realized there is a big problem with the architecture I was planning to design. So the gradients cannot backpropagate through embedding layer too. It means that even if I solve the argmax problem I will still get stuck in the embeddings. I guess thats the reason you proposed to merge the two. I checked again the paper and I think they dont use embeddings at all they just use one-hot vectors (in their first paper they use embeddings without full backpropagation: https://arxiv.org/abs/1707.04879 27).
I this case I will just use something like this:
class ArgMax(torch.autograd.Function):

	@staticmethod
	def forward(ctx, input):
        idx = torch.argmax(input, 1)

        output = torch.zeros_like(input)
        output.scatter_(1, idx, 1)
		
        return output
	

	@staticmethod
	def backward(ctx, grad_output):
        return grad_output

Or simply I will use gumbel softmax: https://pytorch.org/docs/stable/nn.html?highlight=gumbel#torch.nn.functional.gumbel_softmax 279