What is the best way to predict a categorical variable, and then embed it, as input to another net?
My instances are tabular, a mix of categorical and continuous variables. I currently have a siamese net (net1) that uses the instances as input. The categorical variables are integer indices used before an nn.Embedding layer. Once net1 is trained, it is fixed.
Now, I want net0 to predict (output) these tabular instances and feed them as input to the above net, and backprop the loss through the entire network (net0 => net1), even tho net1 will be fixed and not modified.

What seems sensible to me is net0 predicts the one-hot (softmax) representation of the category, and then figure out to change net1 so that nn.Embedding uses a one-hot, not integer index, as the input. However, I don’t believe nn.Embedding can accept a one-hot input (How to make nn.Embedding support one-hot encoded vector input? 12)
net0 predicts the one-hot (softmax) representation of the category, takes the argmax to get the integer index, and net1 looks up the nn.Embedding using the index. However, I believe that argmax is non-differentiable, which means I cannot pass losses back from net1 to net0 Differentiable argmax 15

I could have net0 directly predict the embedding of net1, and in postprogressing use kNN to find the nearest category. I don’t love this because a) it requires the net1 embeddings to be fixed and b) it seems a little inelegant.

What is my best option here to predict a categorical variable and then embed it?