KFrank:

Conceptually, what are the index values in W2 supposed to mean?


I have two networks that are trained together. The networks are in series, i.e. output of one is the input of the other. Upstream network is solving a classification problem for graph edge types [Batch, node, node, class]. What you see is my downstream network: it takes in output of the upstream W, applies Softmax to turn them into probabilities W1, and then I want to turn them into either a one-hot vector or class labels W2 before reshaping and feeding it to the linear layer dec_fc_5.