build_posterior fails after training SNPE on GPU

[Dingel321]

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)`

[michaeldeistler]

Hi there,

thanks for reporting this, it's a bug. I'll fix it right away.

Michael

[michaeldeistler]

The fixup is now available in #730

Merging the branch and making a release will probably take quite some time though, so maybe just checkout that branch