Hey,
I am using the following implementation to apply gumble softmax on my output layer. What loss function should I use with gumble softmax.
I have removed the softmax after applying gumble noise and kept my loss function as CrossEntropyLoss. Is that fine?
Regards