Thanks Thomas,
that really helped me a lot. My embedding layer should be outside of the function based on the architecture that I use. I changed the code slightly like this:
class ArgMax(torch.autograd.Function):

	@staticmethod
	def forward(ctx, input):
		idx = torch.argmax(input, 1)
		ctx._input_shape = input.shape
		ctx._input_dtype = input.dtype
		ctx._input_device = input.device
		ctx.save_for_backward(idx)
		return idx.float()

	@staticmethod
	def backward(ctx, grad_output):
		idx, = ctx.saved_tensors
		grad_input = torch.zeros(ctx._input_shape, device=ctx._input_device, dtype=ctx._input_dtype)
		grad_input.scatter_(1, idx[:, None], grad_output.sum())
		return grad_input


there are gradients propagating back but I’m not 100% sure if they are correct. I need to check them. Also I did not understand what do you mean by  “it doesn’t do gradients on the weight”.