Small modification to WhiteNoiseKernel to support multitask
@@ -15,8 +15,8 @@ class WhiteNoiseKernel(Kernel):
def forward(self, x1, x2):
if self.training:
- return DiagLazyVariable(self.variances.unsqueeze(0))
- 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))
+ return DiagLazyVariable(self.variances.view(-1).unsqueeze(0))
+ elif x1.size(-2) == x2.size(-2) and x1.size(-2) == self.variances.size(0) and torch.equal(x1, x2):
+ return DiagLazyVariable(self.variances.view(-1).unsqueeze(0))
else:
return ZeroLazyVariable(x1.size(-3), x1.size(-2), x2.size(-2))
