Always use output shape instead of input shape for noise shape
@@ -30,13 +30,7 @@ class _GaussianLikelihoodBase(Likelihood):
self.noise_covar = noise_covar
def _shaped_noise_covar(self, base_shape: torch.Size, *params: Any, **kwargs: Any):
- if len(params) > 0:
- # we can infer the shape from the params
- shape = None
- else:
- # here shape[:-1] is the batch shape requested, and shape[-1] is `n`, the number of points
- shape = base_shape
- return self.noise_covar(*params, shape=shape, **kwargs)
+ return self.noise_covar(*params, shape=base_shape, **kwargs)
def expected_log_prob(self, target: Tensor, input: MultivariateNormal, *params: Any, **kwargs: Any) -> Tensor:
mean, variance = input.mean, input.variance
