Repository: sbi-dev/sbi

Issue:
build_posterior fails after training SNPE on GPU
Hi,

I'm training SNPE without an explicit prior on GPU. When I want to build the posterior, I get the following error massage.
If I repeat the training only on the CPU everything works.

`Traceback (most recent call last):
  File "sbi_train_posterior.py", line 68, in <module>
    main(args.train_file, args.posterior_file)
  File "sbi_train_posterior.py", line 45, in main
    posterior = inference.build_posterior(density_estimator)
  File "/home/dingeldein/anaconda3/lib/python3.8/site-packages/sbi/inference/snpe/snpe_base.py", line 420, in build_posterior
    potential_fn, theta_transform = posterior_estimator_based_potential(
  File "/home/dingeldein/anaconda3/lib/python3.8/site-packages/sbi/inference/potentials/posterior_based_potential.py", line 46, in posterior_estimator_based_potential
    theta_transform = mcmc_transform(prior, device=device)
  File "/home/dingeldein/anaconda3/lib/python3.8/site-packages/sbi/utils/sbiutils.py", line 618, in mcmc_transform
    if hasattr(prior, "mean") and hasattr(prior, "stddev"):
  File "/home/dingeldein/anaconda3/lib/python3.8/site-packages/pyro/distributions/empirical.py", line 154, in mean
    return self._weighted_mean(self._samples)
  File "/home/dingeldein/anaconda3/lib/python3.8/site-packages/pyro/distributions/empirical.py", line 136, in _weighted_mean
    return (value * relative_probs).sum(
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!`

My code:

`    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    x_file_name = f'{training_file_name}_x.pt'
    theta_file_name = f'{training_file_name}_theta.pt'

    theta = torch.load(theta_file_name, map_location=torch.device(device))
    x = torch.load(x_file_name, map_location=torch.device(device))

    neural_posterior = posterior_nn(
        model='nsf',
        hidden_features=100,
        num_transforms=6,
        num_bins=10
    )

    inference = SNPE(
        density_estimator=neural_posterior,
        device=device
    )

    inference = inference.append_simulations(theta, x)
    density_estimator = inference.train(
        show_train_summary=True,
        num_atoms=10,
        validation_fraction=0.15,
        training_batch_size=1500,
        learning_rate=0.001,
        stop_after_epochs=30
    )

    posterior = inference.build_posterior(density_estimator)`


Discussion:
Hi there,

thanks for reporting this, it's a bug. I'll fix it right away.

Michael

Resolving patch:
diff --git a/sbi/inference/snpe/snpe_base.py b/sbi/inference/snpe/snpe_base.py
--- a/sbi/inference/snpe/snpe_base.py
+++ b/sbi/inference/snpe/snpe_base.py
@@ -171,8 +171,10 @@ def append_simulations(
                     "run single-round inference with "
                     "`append_simulations(..., proposal=None)`."
                 )
-            theta_prior = self.get_simulations()[0]
-            self._prior = ImproperEmpirical(theta_prior, ones(theta_prior.shape[0]))
+            theta_prior = self.get_simulations()[0].to(self._device)
+            self._prior = ImproperEmpirical(
+                theta_prior, ones(theta_prior.shape[0], device=self._device)
+            )
 
         return self
 


Test patch: