Pyro: [bug] Continuous autoguides do not take device into account

[ahmadsalim]

Issue Description

Many of the AutoContinuous guides do not take into account that the model might be on the GPU instead of the CPU. For example, see the code snippet below, where it fails if I pass device='cuda' to the tensors I create.

It would of course be nice if it instead relied on the device given from the model trace.
Do you have any suggestions?

Thanks!

Environment

• OS and python version: Windows 10, Python 3.7
• PyTorch version: 1.1
• Pyro version: 0.3.3

Code Snippet

import pyro
import pyro.distributions as dist
import torch
from pyro.contrib.autoguide import AutoDiagonalNormal
from pyro.infer import SVI, Trace_ELBO
from pyro.optim import Adam


def model(xs, prior_loc, prior_scale):
    mu = pyro.sample('mu', dist.Normal(prior_loc, prior_scale))
    return pyro.sample('obs', dist.Normal(mu, prior_scale), obs=xs)


def train():
    xs = torch.tensor([5, 5.5, 5.6, 5.8, 6.3], device='cuda')
    prior_loc = torch.tensor(0.0, device='cuda')
    prior_scale = torch.tensor(1.0, device='cuda')
    guide = AutoDiagonalNormal(model)
    optim = Adam({'lr': 1e-3})
    loss = Trace_ELBO()
    svi = SVI(model, guide, optim, loss)
    svi.step(xs, prior_loc, prior_scale)

if __name__ == "__main__":
  train()

Error:

Traceback (most recent call last):
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\IPython\core\interactiveshell.py", line 3265, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-6-68d49e4377f9>", line 26, in <module>
    train()
  File "<ipython-input-6-68d49e4377f9>", line 23, in train
    svi.step(xs, prior_loc, prior_scale)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\pyro\infer\svi.py", line 99, in step
    loss = self.loss_and_grads(self.model, self.guide, *args, **kwargs)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\pyro\infer\trace_elbo.py", line 136, in loss_and_grads
    surrogate_loss_particle.backward(retain_graph=self.retain_graph)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\torch\tensor.py", line 107, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "C:\tools\miniconda3\envs\scanner-pyro\lib\site-packages\torch\autograd\__init__.py", line 93, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type torch.cuda.FloatTensor but got torch.FloatTensor

[jpchen]

this seems to come from the initialization of the autoguide parameters. i will make a fix for this but in the meantime you can set the default tensor type via torch.set_default_tensor_type('torch.cuda.FloatTensor') and it should work without the device kwarg to each tensor.

[ahmadsalim]

Awesome, thank you very much! 😄

[fritzo]

@jpchen I am refactoring autoguide initialization in #1849, so I will take care of this in that PR.