Thanks! I think the problem stems from the following aspect:" during the sampling process of NUTS the newly accepted samples of parameters are in closed form, which can not be passed back to the network, and the code need to recompute GMM every time the newly accepted samples of parameters are founded. What I need is something like BP algorithm, the updation can be obtained with the iteration"