Update: the code is here:

https://github.com/jerryqhyu/distill_bayes_net