def METHOD_NAME(self, xs, ys=None):
    # register this pytorch module and all of its sub-modules with pyro
    pyro.module("generation_net", self)
    batch_size = xs.shape[0]
    with pyro.plate("data"):
        # Prior network uses the baseline predictions as initial guess.
        # This is the generative process with recurrent connection
        with torch.no_grad():
            # this ensures the training process does not change the
            # baseline network
            y_hat = self.baseline_net(xs).view(xs.shape)
        # sample the handwriting style from the prior distribution, which is
        # modulated by the input xs.
        prior_loc, prior_scale = self.prior_net(xs, y_hat)
        zs = pyro.sample("z", dist.Normal(prior_loc, prior_scale).to_event(1))
        # the output y is generated from the distribution pθ(y|x, z)
        loc = self.generation_net(zs)
        if ys is not None:
            # In training, we will only sample in the masked image
            mask_loc = loc[(xs == -1).view(-1, 784)].view(batch_size, -1)
            mask_ys = ys[xs == -1].view(batch_size, -1)
            pyro.sample(
                "y",
                dist.Bernoulli(mask_loc, validate_args=False).to_event(1),
                obs=mask_ys,
            )
        else:
            # In testing, no need to sample: the output is already a
            # probability in [0, 1] range, which better represent pixel
            # values considering grayscale. If we sample, we will force
            # each pixel to be  either 0 or 1, killing the grayscale
            pyro.deterministic("y", loc.detach())
        # return the loc so we can visualize it later
        return loc