There are two ways of going about this. In one, you can implement a new
GANInference subclass while perturbing all the gradient updates with SGHMC noise terms. One logistical difficulty is figuring out how to store these new parameters (which is an "MCMC sample") before it gets overwritten on the next iteration.
In another, you can implement a new
SGHMC subclass. You would rewrite some of the logic around the log prob calculation so that it can alternate updates between two sets of parameters.