Likelihood passes in args/kwargs to expected_log_prob

Author: gpleiss

gpytorch/likelihoods/likelihood.py

@@ -45,7 +45,7 @@ def expected_log_prob(
         self, observations: Tensor, function_dist: MultivariateNormal, *args: Any, **kwargs: Any
     ) -> Tensor:
         likelihood_samples = self._draw_likelihood_samples(function_dist, *args, **kwargs)
-        res = likelihood_samples.log_prob(observations).mean(dim=0)
+        res = likelihood_samples.log_prob(observations, *args, **kwargs).mean(dim=0)
         return res
 
     @abstractmethod
@@ -410,7 +410,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
     def expected_log_prob(
         self, observations: Tensor, function_dist: MultivariateNormal, *args: Any, **kwargs: Any
     ) -> Tensor:
-        log_prob_lambda = lambda function_samples: self.forward(function_samples).log_prob(observations)
+        log_prob_lambda = lambda function_samples: self.forward(function_samples, *args, **kwargs).log_prob(
+            observations
+        )
         log_prob = self.quadrature(log_prob_lambda, function_dist)
         return log_prob