Subsampling with both data-loaders and subsampling?

1

Hi, I have a question on how to subsample when I have 2 parameters that I want to subsample. The problem is that I am using ammortized inference on 1 of them, but I do not want to use ammortized inference on the other parameter. I am already using a dataloader so the data that feed in is already subsampled.

I am trying to subsample variable z2, which is not ammortized
Example code:

def model(data, idx): 
   with pyro.plate('batch', 1024):
    z1 = pyro.sample("z1", dist.Normal(...))
    z2 = ?
def guide(data, idx): 
   with pyro.plate('batch', 1024):
     z1_loc, z1_scale = encoder(data)
     z1 = pyro.sample("z1", dist.Normal(z1_loc, z1_scale))
     z2 = ?
2

if i understand your question right it looks like you need to pass your idx to plate

pyro.plate('batch', 1024, subsample=idx):

see docs: Primitives — Pyro documentation