fix: infer device with index from density estimator.

sbi-dev · Jun 25, 2024 · e798a03 · e798a03
1 parent 6f61662
commit e798a03
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/sbi/inference/snle/snle_base.py b/sbi/inference/snle/snle_base.py
@@ -321,7 +321,7 @@ def build_posterior(
         else:
             likelihood_estimator = density_estimator
             # Otherwise, infer it from the device of the net parameters.
-            device = next(density_estimator.parameters()).device.type
+            device = str(next(density_estimator.parameters()).device)
 
         potential_fn, theta_transform = likelihood_estimator_based_potential(
             likelihood_estimator=likelihood_estimator,

diff --git a/sbi/inference/snpe/snpe_base.py b/sbi/inference/snpe/snpe_base.py
@@ -500,7 +500,7 @@ def build_posterior(
         else:
             posterior_estimator = density_estimator
             # Otherwise, infer it from the device of the net parameters.
-            device = next(density_estimator.parameters()).device.type
+            device = str(next(density_estimator.parameters()).device)
 
         potential_fn, theta_transform = posterior_estimator_based_potential(
             posterior_estimator=posterior_estimator,

diff --git a/sbi/inference/snre/snre_base.py b/sbi/inference/snre/snre_base.py
@@ -376,7 +376,7 @@ def build_posterior(
         else:
             ratio_estimator = density_estimator
             # Otherwise, infer it from the device of the net parameters.
-            device = next(density_estimator.parameters()).device.type
+            device = str(next(density_estimator.parameters()).device)
 
         potential_fn, theta_transform = ratio_estimator_based_potential(
             ratio_estimator=ratio_estimator,