You’d need to return a zero vector with the argmax entry set to a gradient, this can be achieved by zeros followed by assignment, scatter, or a similar function.
Note that what you do makes very limited sense mathematically (because you trade output for indexing) and argmax is an integral tensor which cannot have a gradient, so you’d need to cast it to float before doing much with it, …
Best of luck
Thomas