allow gpu when using an empirical prior for SNPE

sbi-dev · Aug 24, 2022 · 893ac43 · 893ac43
1 parent 6fed0f4
commit 893ac43
Showing 1 changed file with 4 additions and 2 deletions.
diff --git 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