Fix ZeroLazyVariable size returned by white noise kernel
@@ -14,4 +14,4 @@ class WhiteNoiseKernel(Kernel):
elif x1.size(-2) == x2.size(-2) and x1.size(-2) == self.variances.size(-1) and torch.equal(x1, x2):
return DiagLazyVariable(self.variances.unsqueeze(0))
else:
- return ZeroLazyVariable(x1.size(-3), x1.size(-2), x1.size(-1), device=x1.device)
+ return ZeroLazyVariable(x1.size(-3), x1.size(-2), x2.size(-2), device=x1.device)
