Save only a subset of latent variables in MCMC?

Hi @fehiepsi and @neerajprad, I’d like to perform a long MCMC inference in Pyro on GPU. I have two latent variables: one x of size about 2000, one y of size about 200000. My model’s geometry looks good and NUTS tunes to step size of about 0.03 which I believe is reasonable.

My problem is that I care only about the smaller variable x and would like to discard samples from the larger variable y during inference (to save memory) and avoid reporting on y in mcmc.summary(). Is there a way to ignore some latent variables when saving samples, e.g. some way to do this with hook_fn()? Thanks in advance!

This is solved by @fritzo in this PR. :smiley: