Repository: pyro-ppl/pyro

Issue:
`enum_discrete` Tensor location bug
Hi,

When using `enum_discrete = True` to reduce variance in my SVI, Pyro complains that some internal tensors are not in the same location (GPU/CPU), whereas turning it off makes it work fine:

```python-traceback
---------------------------------------------------------------------------
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'
```

I suspect this is a bug. Any tips?

-Reinier

PS The bug arises here:
```python
        for inputs, targets in train_loader:

            # Put batch on GPU
            input_batch = Variable(inputs.cuda(device=CUDA_ID, async=use_async))  # Pinned memory with async transfer
            target_batch = Variable(targets.cuda(device=CUDA_ID, async=use_async))

            # Train and score
            train_loss += svi.step(input_batch, target_batch)

        # Store history
        train_losses += [train_loss / n_train_samples]
```
Where `CUDA_ID` is some device id and async is currently `False`.


Discussion:
Thanks for pinpointing this. I think we have a ton of these placement bugs. I'll take a look.

Resolving patch:
diff --git a/pyro/distributions/util.py b/pyro/distributions/util.py
--- a/pyro/distributions/util.py
+++ b/pyro/distributions/util.py
@@ -151,7 +151,7 @@ def torch_multinomial(input, num_samples, replacement=False):
     Does not support keyword argument `out`.
     """
     if input.is_cuda:
-        return torch_multinomial(input.cpu(), num_samples, replacement).cuda()
+        return torch.multinomial(input.cpu(), num_samples, replacement).cuda(input.get_device())
     else:
         return torch.multinomial(input, num_samples, replacement)
 
diff --git a/pyro/infer/enum.py b/pyro/infer/enum.py
--- a/pyro/infer/enum.py
+++ b/pyro/infer/enum.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 import functools
+import math
 
 import torch
 from torch.autograd import Variable
@@ -44,9 +45,8 @@ def iter_discrete_traces(graph_type, fn, *args, **kwargs):
 
         # 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())
+        if isinstance(log_pdf, Variable):
+            scale = torch.exp(log_pdf.detach())
+        else:
+            scale = math.exp(log_pdf)
         yield scale, full_trace
diff --git a/pyro/infer/trace_elbo.py b/pyro/infer/trace_elbo.py
--- a/pyro/infer/trace_elbo.py
+++ b/pyro/infer/trace_elbo.py
@@ -4,6 +4,7 @@
 import warnings
 
 import numpy as np
+from torch.autograd import Variable
 
 import pyro
 import pyro.poutine as poutine
@@ -100,7 +101,10 @@ def loss(self, model, guide, *args, **kwargs):
         for weight, model_trace, guide_trace, log_r in self._get_traces(model, guide, *args, **kwargs):
             elbo_particle = weight * 0
 
-            log_pdf = "batch_log_pdf" if (self.enum_discrete and weight.size(0) > 1) else "log_pdf"
+            if (self.enum_discrete and isinstance(weight, Variable) and weight.size(0) > 1):
+                log_pdf = "batch_log_pdf"
+            else:
+                log_pdf = "log_pdf"
             for name in model_trace.nodes.keys():
                 if model_trace.nodes[name]["type"] == "sample":
                     if model_trace.nodes[name]["is_observed"]:
@@ -137,7 +141,10 @@ def loss_and_grads(self, model, guide, *args, **kwargs):
             elbo_particle = weight * 0
             surrogate_elbo_particle = weight * 0
             # compute elbo and surrogate elbo
-            log_pdf = "batch_log_pdf" if (self.enum_discrete and weight.size(0) > 1) else "log_pdf"
+            if (self.enum_discrete and isinstance(weight, Variable) and weight.size(0) > 1):
+                log_pdf = "batch_log_pdf"
+            else:
+                log_pdf = "log_pdf"
             for name, model_site in model_trace.nodes.items():
                 if model_site["type"] == "sample":
                     if model_site["is_observed"]:


Test patch: