fix: device mismatch in NPSE marginal mean/std computation (#1707)

* Fix device mismatch in NPSE marginal mean/std computation

* Apply pre-commit formatting to score_estimator fix

* Inline device assignment for t_tensor as suggested by reviewer

Co-authored-by: Jan Teusen (né Boelts) <janfb@users.noreply.github.com>

---------

Co-authored-by: Jan Teusen (né Boelts) <janfb@users.noreply.github.com>

Author: satwiksps

sbi/neural_nets/estimators/score_estimator.py

@@ -118,8 +118,10 @@ def __init__(
 
         # Now that input_shape and mean_0, std_0 is set, we can compute the proper mean
         # and std for the "base" distribution.
-        mean_t = self.approx_marginal_mean(torch.tensor([t_max]))
-        std_t = self.approx_marginal_std(torch.tensor([t_max]))
+        # Create t on the correct device to avoid CPU/GPU mismatch
+        t_tensor = torch.as_tensor([t_max], device=self.mean_0.device)
+        mean_t = self.approx_marginal_mean(t_tensor)
+        std_t = self.approx_marginal_std(t_tensor)
         mean_t = torch.broadcast_to(mean_t, (1, *input_shape))
         std_t = torch.broadcast_to(std_t, (1, *input_shape))