# Buggy code
def iter_discrete_traces(graph_type, fn, *args, **kwargs):
    """
    Iterate over all discrete choices of a stochastic function.

    When sampling continuous random variables, this behaves like `fn`.
    When sampling discrete random variables, this iterates over all choices.

    This yields `(scale, trace)` pairs, where `scale` is the probability of the
    discrete choices made in the `trace`.

    :param str graph_type: The type of the graph, e.g. "flat" or "dense".
    :param callable fn: A stochastic function.
    :returns: An iterator over (scale, trace) pairs.
    """
    queue = LifoQueue()
    queue.put(Trace())
    while not queue.empty():
        partial_trace = queue.get()
        escape_fn = functools.partial(util.discrete_escape, partial_trace)
        traced_fn = poutine.trace(
            poutine.escape(poutine.replay(fn, partial_trace), escape_fn),
            graph_type=graph_type,
        )
        try:
            full_trace = traced_fn.get_trace(*args, **kwargs)
        except util.NonlocalExit as e:
            for extended_trace in util.enum_extend(traced_fn.trace.copy(), e.site):
                queue.put(extended_trace)
            continue

        # Scale trace by probability of discrete choices.
        log_pdf = full_trace.batch_log_pdf(site_filter=site_is_discrete)
        if isinstance(log_pdf, float):
            log_pdf = torch.Tensor([log_pdf])
        if isinstance(log_pdf, torch.Tensor):
            log_pdf = Variable(log_pdf)
        scale = torch.exp(log_pdf.detach())
        yield scale, full_trace


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-31-93fd337a0d5c> in <module>()
----> 1 train()

<ipython-input-30-98110a9a9720> in train()
33
34             # Train and score
---> 35             train_loss += svi.step(input_batch, target_batch)
36
37         # Store history

/home/reinier/anaconda2/lib/python2.7/site-packages/pyro/infer/svi.pyc in step(self, *args, **kwargs)
96         """
97         # get loss and compute gradients
---> 98         loss = self.loss_and_grads(self.model, self.guide, *args, **kwargs)
99
100         # get active params

/home/reinier/anaconda2/lib/python2.7/site-packages/pyro/infer/elbo.pyc in loss_and_grads(self, model, guide, *args, **kwargs)
63         :rtype: float
64         """
---> 65         return self.which_elbo.loss_and_grads(model, guide, *args, **kwargs)

/home/reinier/anaconda2/lib/python2.7/site-packages/pyro/infer/trace_elbo.pyc in loss_and_grads(self, model, guide, *args, **kwargs)
149                         guide_site = guide_trace.nodes[name]
150                         lp_lq = model_site[log_pdf] - guide_site[log_pdf]
--> 151                         elbo_particle += lp_lq
152                         if guide_site["fn"].reparameterized:
153                             surrogate_elbo_particle += lp_lq

RuntimeError: Expected object of type Variable[torch.FloatTensor] but found type Variable[torch.cuda.FloatTensor] for argument #1 'other'

# Fixed code
def iter_discrete_traces(graph_type, fn, *args, **kwargs):
    """
    Iterate over all discrete choices of a stochastic function.

    When sampling continuous random variables, this behaves like `fn`.
    When sampling discrete random variables, this iterates over all choices.

    This yields `(scale, trace)` pairs, where `scale` is the probability of the
    discrete choices made in the `trace`.

    :param str graph_type: The type of the graph, e.g. "flat" or "dense".
    :param callable fn: A stochastic function.
    :returns: An iterator over (scale, trace) pairs.
    """
    queue = LifoQueue()
    queue.put(Trace())
    while not queue.empty():
        partial_trace = queue.get()
        escape_fn = functools.partial(util.discrete_escape, partial_trace)
        traced_fn = poutine.trace(
            poutine.escape(poutine.replay(fn, partial_trace), escape_fn),
            graph_type=graph_type,
        )
        try:
            full_trace = traced_fn.get_trace(*args, **kwargs)
        except util.NonlocalExit as e:
            for extended_trace in util.enum_extend(traced_fn.trace.copy(), e.site):
                queue.put(extended_trace)
            continue

        # Scale trace by probability of discrete choices.
        log_pdf = full_trace.batch_log_pdf(site_filter=site_is_discrete)
        if isinstance(log_pdf, Variable):
            scale = torch.exp(log_pdf.detach())
        else:
            scale = math.exp(log_pdf)
        yield scale, full_trace