GPU compatibility

sbi/samplers/rejection/rejection.py

@@ -258,6 +258,10 @@ def accept_reject_sample(
         proposal_sampling_kwargs = {}
 
     num_remaining = num_samples
+
+    # NOTE: We might want to change this to a more general approach in the future.
+    # Currently limited to a single "batch_dim" for the condition.
+    # But this would require giving the method the condition_shape explicitly...
     if "condition" in proposal_sampling_kwargs:
         num_xos = proposal_sampling_kwargs["condition"].shape[0]
     else:
@@ -295,7 +299,7 @@ def accept_reject_sample(
         # samples will be of shape(sampling_batch_size,*batch_shape, *event_shape)
         # and hence work in dim = 0.
         num_accepted = are_accepted.sum(dim=0)
-        num_sampled_total += num_accepted
+        num_sampled_total += num_accepted.to(num_sampled_total.device)
         num_samples_possible += sampling_batch_size
         min_num_accepted = num_accepted.min().item()
         num_remaining -= min_num_accepted
@@ -358,4 +362,4 @@ def accept_reject_sample(
         samples.shape[0] == num_samples
     ), "Number of accepted samples must match required samples."
 
-    return samples, as_tensor(min_acceptance_rate)
+    return samples, acceptance_rate