def METHOD_NAME(self, grads, batch, *args, **kwargs):
    """
    :returns: an estimate of the loss (expectation over p(x, y) of
        -log q(x, y) ) - where p is the model and q is the guide
    :rtype: float
    If a batch is provided, the loss is estimated using these traces
    Otherwise, a fresh batch is generated from the model.
    If grads is True, will also call `backward` on loss.
    `args` and `kwargs` are passed to the model and guide.
    """
    if batch is None:
        batch = (
            self._sample_from_joint(*args, **kwargs)
            for _ in range(self.training_batch_size)
        )
        batch_size = self.training_batch_size
    else:
        batch_size = len(batch)
    loss = 0
    for model_trace in batch:
        with poutine.trace(param_only=True) as particle_param_capture:
            guide_trace = self._get_matched_trace(model_trace, *args, **kwargs)
        particle_loss = self._differentiable_loss_particle(guide_trace)
        particle_loss /= batch_size
        if grads:
            guide_params = set(
                site["value"].unconstrained()
                for site in particle_param_capture.trace.nodes.values()
            )
            guide_grads = torch.autograd.grad(
                particle_loss, guide_params, allow_unused=True
            )
            for guide_grad, guide_param in zip(guide_grads, guide_params):
                if guide_grad is None:
                    continue
                guide_param.grad = (
                    guide_grad
                    if guide_param.grad is None
                    else guide_param.grad + guide_grad
                )
        loss += torch_item(particle_loss)
    warn_if_nan(loss, "loss")
    return loss