Straight through gradient estimator

I am wondering whether straight through gradient estimator for categorical variables is implemented. I want to see if I can implement the vq-vae architecture with Pyro. Thanks

We have an implementation for relaxed categorical straight through in pyro.distributions.relaxed_straight_through on our current dev branch. You should be able to use that for your purpose (and it will be great if you could contribute it as a Pyro example!). Let us know if you face any issues.

That’s great. I have modified the vae example with convolutional layers. Next up is the vq implementation. I might also leave this for students as part of a capstone project.