How to apply Gumbel softmax to node features?

[gooaah]

Hi everyone. I want to make continuous node features become one-hot before pooling but keep it differentiable (Gumbel softmax).

Is it implemented in the current version? If not, how to do it?

[rusty1s]

I think you can use the regular gumbel_softmax implementation of PyTorch for this.

[gooaah]

You are right, thanks!