Thank you for explaining your use case! I do think that it is likely easiest to integrate argmax and embedding into a single autograd Function.
To use it I’d probably adapt the “later model” to (optionally) take inputs before argmax and do the integrated argmax + embedding (my personal style would be to pass a flag on construction whether that should be the training mode and then test that and self.training during forward).
Best regards
Thomas