Fix ScaleKernel diag
@@ -92,7 +92,7 @@ class ScaleKernel(Kernel):
if batch_dims == (0, 2) and outputscales.numel() > 1:
outputscales = outputscales.unsqueeze(1).repeat(1, x1.size(-1)).view(-1)
- orig_output = self.base_kernel(x1, x2, batch_dims=batch_dims, **params)
+ orig_output = self.base_kernel(x1, x2, diag=diag, batch_dims=batch_dims, **params)
if torch.is_tensor(orig_output):
outputscales = outputscales.view(-1, *([1] * (orig_output.dim() - 1)))
return orig_output.mul(outputscales)
