I need to put argmax in the middle of my network and thus I need it to be differentiable using straight-through estimator, thats: during the forward I want to do the usual argmax and during the backward, as argmax is not differentiable, I would like to pass the incoming gradient instead of 0 gradients. This is what I came up with:
class ArgMax(torch.autograd.Function):

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


	@staticmethod
	def backward(ctx, grad_output):
		grad_input = grad_output.clone()
		return grad_input

However I get this error:
RuntimeError: RuntimeError: Function ArgMaxBackward returned an invalid gradient at index 0 - expected shape [1, 10] but got [1]
The problem is that the input to the forward function is a tensor of [1, 10] but because of the armax, the its output becomes a  tensor of size [1]. I dont know the mechanics of autograd very well. Is there is solution for this problem?