squeeze last dim of ratios

sbi-dev · Apr 9, 2024 · 0e62d37 · 0e62d37
1 parent 3b8afae
commit 0e62d37
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sbi/neural_nets/ratio_estimators.py b/sbi/neural_nets/ratio_estimators.py
@@ -64,6 +64,7 @@ def unnormalized_log_ratio(self, theta: Tensor, x: Tensor, **kwargs) -> Tensor:
 
         Returns:
             Sample-wise unnormalized log ratios.
+            Just like log_prob, the last dimension should be squeezed.
         """
 
         raise NotImplementedError
@@ -104,4 +105,4 @@ def unnormalized_log_ratio(self, theta: Tensor, x: Tensor) -> Tensor:
         z = self.combine_embedded_theta_and_x(
             self.embedding_net_theta(theta), self.embedding_net_x(x)
         )
-        return self.net(z)
+        return self.net(z).squeeze(-1)