Hi Thomas!
That was clear, thanks! The solution is to use the Gumbel-Softmax and to make sure the gradients remain for the category probabilities, I did the following (dot product mm):
z, probs = g_s_net(input, temp)
sample = torch.mm(z, torch.FloatTensor(some_categorical_list))
loss = probs * z

Here z is the categorical one hot tensor and multiplying it with the probs will give the probs for the corresponding class.